Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2f6bce1632 |
No files matched your search
@@ -30,6 +30,14 @@ steps:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: "Reliability Test - voxcpm2"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- group: "Reliability Test - Invalid parameters Test"
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 96 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 112 KiB |
@@ -38,12 +38,11 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
|
||||
| `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
|
||||
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. |
|
||||
| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. |
|
||||
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
|
||||
| `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
|
||||
| `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
|
||||
| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
|
||||
| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
|
||||
| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
|
||||
| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
|
||||
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
|
||||
| `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
|
||||
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
|
||||
|
||||
@@ -184,6 +184,9 @@ Failed jobs: 1/2
|
||||
# Reliability: HunyuanImage DiT (H100 × 4)
|
||||
pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
|
||||
# Reliability: VoxCPM2 (L4)
|
||||
pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
|
||||
```
|
||||
|
||||
The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml).
|
||||
|
||||
@@ -562,6 +562,7 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
|
||||
- `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
|
||||
- `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
|
||||
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
|
||||
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
|
||||
- `README.md` — Short local run commands for this directory.
|
||||
@@ -661,12 +662,13 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
|
||||
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
|
||||
| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
|
||||
| `README.md` | Minimal run / collect examples. |
|
||||
|
||||
##### Parametrization and markers
|
||||
|
||||
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100).
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
|
||||
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
|
||||
|
||||
##### CI trigger
|
||||
@@ -678,11 +680,13 @@ Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model sui
|
||||
| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
|
||||
| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
```
|
||||
|
||||
##### Local commands
|
||||
@@ -692,6 +696,7 @@ pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
```
|
||||
|
||||
##### Adding a new model suite
|
||||
@@ -704,4 +709,4 @@ pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
|
||||
@@ -10,30 +10,6 @@ For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
|
||||
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||
|
||||
## Feature compatibility
|
||||
|
||||
Host-storage optimization and runtime compatibility are separate decisions.
|
||||
When direct checkpoint mmap is unavailable, DLO can still use tensors produced
|
||||
by the ordinary loader. "Compatibility path" below means that fallback is
|
||||
implemented but has less end-to-end coverage than the primary path.
|
||||
|
||||
Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
|
||||
|
||||
| Feature | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
|
||||
| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
|
||||
| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
|
||||
| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
|
||||
| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
|
||||
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
|
||||
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
|
||||
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
|
||||
|
||||
See [Parallelism compatibility](#parallelism-compatibility) and
|
||||
[Request and loading constraints](#request-and-loading-constraints) for the
|
||||
detailed contracts and validation boundaries.
|
||||
|
||||
## Status
|
||||
|
||||
DLO is implemented for multi-device diffusion execution. The default
|
||||
@@ -205,11 +181,8 @@ requirements.
|
||||
|
||||
Direct checkpoint mmap can back either transfer path. It is currently limited
|
||||
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
||||
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
|
||||
ordinary loader finalizes their runtime weights and scales; DLO then shards and
|
||||
reconstructs those tensors with their recorded layouts. Other online methods
|
||||
must use `--dlo-no-use-allgather` or disable online quantization until their
|
||||
runtime layouts are validated.
|
||||
ordinary loader. Online quantization remains incompatible with DLO AllGather;
|
||||
use `--dlo-no-use-allgather` or disable online quantization.
|
||||
|
||||
A normalized runtime mmap cache, built through the ordinary loader, is the
|
||||
proposed general mechanism for sharing transformed TP or quantized layouts.
|
||||
@@ -224,8 +197,6 @@ Current source-level validation includes:
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
|
||||
loaders, missing keys, and shape/dtype mismatches;
|
||||
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
|
||||
sharding of finalized weights and scales;
|
||||
- exact loader-to-backend plan transfer and ordinary-loader fallback;
|
||||
- rank-local mmap source retention, bounded two-slot staging, and adapter
|
||||
transforms without parameter-side flags;
|
||||
|
||||
@@ -190,15 +190,12 @@ configured parallelism strategy:
|
||||
- `WeightsMapper` transformations must be applied before matching ignored-layer
|
||||
lists or component prefixes.
|
||||
|
||||
Distributed layerwise offload has an additional boundary: direct checkpoint
|
||||
mmap cannot prepare online-quantized tensors because it bypasses the ordinary
|
||||
loader that creates their runtime weights and scales. Per-tensor online FP8
|
||||
linears can still use DLO AllGather through the ordinary loader: each rank first
|
||||
finalizes the FP8 tensors, then DLO shards those runtime tensors for
|
||||
reconstruction during prefetch. Other online methods remain limited to
|
||||
`--dlo-no-use-allgather` until their packing and scale layouts are validated.
|
||||
See the [DLO design document](offloader/distributed_layerwise_offload.md) for
|
||||
the full compatibility matrix.
|
||||
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](offloader/distributed_layerwise_offload.md)
|
||||
for the full compatibility matrix.
|
||||
|
||||
## Adding a quantization backend
|
||||
|
||||
|
||||
@@ -18,8 +18,8 @@ validation_paths:
|
||||
- tests/config/**
|
||||
upstream_refs:
|
||||
- vllm.config
|
||||
last_reviewed: 2026-08-11
|
||||
last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
---
|
||||
|
||||
# vLLM-Omni configuration
|
||||
@@ -39,66 +39,11 @@ The PR review proposal to capture environment-derived values in configuration
|
||||
objects at initialization is an open design question, not a current
|
||||
invariant.
|
||||
|
||||
## Current structured-stage reuse behavior
|
||||
|
||||
The current RFC #4021 implementation uses upstream vLLM config classes as the
|
||||
schema source for AR and generation stages:
|
||||
|
||||
- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`.
|
||||
- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`.
|
||||
- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`.
|
||||
- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`.
|
||||
|
||||
This inheritance preserves the existing load, cache, scheduler, and parallel
|
||||
concern boundaries while reducing duplicate downstream declarations. It does
|
||||
not, by itself, make every inherited field an effective vLLM-Omni runtime
|
||||
option. Construction and runtime consumption are separate surfaces:
|
||||
|
||||
| Surface | Current behavior |
|
||||
| --- | --- |
|
||||
| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. |
|
||||
| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. |
|
||||
| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. |
|
||||
| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. |
|
||||
|
||||
The generated projection maps include known upstream naming differences such
|
||||
as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and
|
||||
`data_parallel_master_ip` to `data_parallel_address`. A field added to both an
|
||||
upstream concern config and `EngineArgs` therefore enters the AR/generation
|
||||
projection without requiring a second downstream allowlist. Explicit inherited
|
||||
fields that have no projection raise an error rather than being silently
|
||||
dropped.
|
||||
|
||||
Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`,
|
||||
cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns
|
||||
`distributed_executor_backend` and `worker_cls`; vLLM's private API-process
|
||||
fields are terminal internals. These exclusions prevent one input from being
|
||||
constructed or projected through two config concerns. Effective-engine-argument
|
||||
tests cover both the dynamic mapping and this exclusion boundary.
|
||||
|
||||
`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields:
|
||||
`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML
|
||||
inputs and complete upstream value objects. Mapping inputs are validated and
|
||||
materialized as their concrete upstream types when the structured config is
|
||||
constructed, so downstream consumers see one resolved representation.
|
||||
Prebuilt upstream objects retain their type across the structured projection
|
||||
boundary.
|
||||
|
||||
Quantization remains on its existing Omni-owned transport contract. This
|
||||
reuse step does not adopt upstream `QuantizationConfigArgs`, because the
|
||||
current engine-specific quantization split needs to be resolved separately.
|
||||
|
||||
Diffusion keeps its existing engine-field ownership and projection surface.
|
||||
Shared Python inheritance must not be treated as evidence that an LLM-only
|
||||
parallel input is effective for a diffusion stage.
|
||||
|
||||
## Safe-change guide while deferred
|
||||
|
||||
Changes should trace every affected producer to its runtime consumer and test
|
||||
the structured, legacy, CLI, and deployment paths that are actually supported.
|
||||
For inherited vLLM fields, tests should separately cover constructor behavior
|
||||
and effective `EngineArgs` projection. Do not infer a stable contract from
|
||||
this placeholder page.
|
||||
Do not infer a stable contract from this placeholder page.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
|
||||
@@ -39,7 +39,6 @@ th {
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
|
||||
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
|
||||
@@ -79,7 +78,6 @@ th {
|
||||
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
|
||||
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
|
||||
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
|
||||
@@ -138,11 +138,8 @@ must enter each collective.
|
||||
but it provides no shared-mmap host-memory benefit or guarantee.
|
||||
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
|
||||
AllGather has limited end-to-end validation.
|
||||
- Per-tensor online FP8 linears use the ordinary loader and can run with either
|
||||
DLO transfer path. With AllGather, every rank temporarily materializes the
|
||||
complete FP8 model in host memory before DLO retains only its shard. Other
|
||||
online quantization methods require no-AllGather until their runtime layouts
|
||||
are validated.
|
||||
- Online quantization uses the ordinary loader with no-AllGather. It remains
|
||||
incompatible with DLO AllGather.
|
||||
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
||||
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
||||
- DP concurrency requires an explicit, identical inference-step count.
|
||||
|
||||
@@ -152,8 +152,7 @@ In `DiffusionParallelConfig`:
|
||||
**Constraints:**
|
||||
|
||||
- `hsdp_replicate_size × hsdp_shard_size == world_size`
|
||||
- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
|
||||
Pipeline Parallelism, or Expert Parallelism.
|
||||
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -168,14 +168,6 @@ The following tables show which models support each feature:
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
|
||||
| **SANA-WM** | ❌ | ❌ | ❌<sup>5</sup> | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
|
||||
> Notes:
|
||||
> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta
|
||||
> recurrence carries state across frames, so a rank cannot denoise a slice of
|
||||
> the token sequence in isolation. Doing so would need a distributed scan or
|
||||
> an all-gather before every GDN block. The remaining ❌ columns are simply
|
||||
> unvalidated on this model, not known-broken.
|
||||
|
||||
> **Step execution note:** Helios supports single-request step execution only;
|
||||
> use `max_num_seqs=1`.
|
||||
@@ -183,7 +175,7 @@ The following tables show which models support each feature:
|
||||
**Frame Interpolation Support**
|
||||
|
||||
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
|
||||
### AudioGen
|
||||
|
||||
|
||||
@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
|
||||
|
||||
## Source of Conclusions
|
||||
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
|
||||
@@ -179,24 +179,6 @@ python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
```
|
||||
|
||||
### Distributed layerwise offload
|
||||
|
||||
Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
|
||||
|
||||
```bash
|
||||
vllm serve <your-model> \
|
||||
--omni \
|
||||
--quantization fp8 \
|
||||
--enable-distributed-layerwise-offload
|
||||
```
|
||||
|
||||
Direct checkpoint mmap is not used for this combination. Each rank runs the
|
||||
ordinary loader to create the final FP8 weights and scales before DLO shards
|
||||
them, so startup temporarily materializes the complete FP8 model in host memory
|
||||
on every rank. Runtime DLO host residency is sharded normally. Use
|
||||
`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract
|
||||
does not fit the workload.
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|
||||
@@ -53,6 +53,7 @@ from typing import Any
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
@@ -138,6 +139,12 @@ async def main():
|
||||
|
||||
# ---- Torch setup ----
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
|
||||
# ---- Cache Config ----
|
||||
if args.cache_backend == "cache_dit":
|
||||
@@ -164,10 +171,7 @@ async def main():
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
diffusion_load_format="dummy",
|
||||
|
||||
@@ -154,6 +154,8 @@ def _positive_finite(value: float, flag: str) -> float:
|
||||
def build_omni_kwargs(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
*,
|
||||
parallel_config: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Build Omni configuration without overriding checkpoint class discovery."""
|
||||
|
||||
@@ -165,7 +167,7 @@ def build_omni_kwargs(
|
||||
return {
|
||||
"model": model,
|
||||
"flow_shift": flow_shift,
|
||||
"tensor_parallel_size": args.tensor_parallel_size,
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"model_config": {"lingbot_action_root": str(paths.action_root)},
|
||||
}
|
||||
@@ -260,10 +262,12 @@ def main(argv: Sequence[str] | None = None) -> Path:
|
||||
# helper tests work on machines without the native vLLM runtime.
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni_kwargs = build_omni_kwargs(args, paths)
|
||||
parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
|
||||
omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
|
||||
omni = Omni(**omni_kwargs)
|
||||
try:
|
||||
outputs = omni.generate(
|
||||
|
||||
@@ -125,6 +125,7 @@ async def run(argv: Sequence[str] | None = None) -> Path:
|
||||
# require a CUDA-enabled vLLM installation.
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
@@ -140,7 +141,7 @@ async def run(argv: Sequence[str] | None = None) -> Path:
|
||||
model=args.model,
|
||||
engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
|
||||
enforce_eager=args.enforce_eager,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
|
||||
max_num_seqs=1,
|
||||
model_config={
|
||||
"ar_diffusion_height": args.height,
|
||||
|
||||
@@ -96,6 +96,7 @@ from typing import Any
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
@@ -451,6 +452,14 @@ def main():
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
)
|
||||
|
||||
# Configure cache based on backend type
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
@@ -481,11 +490,7 @@ def main():
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
parallel_config=parallel_config,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
profiler_config=args.profiler_config,
|
||||
|
||||
@@ -56,6 +56,7 @@ import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -502,6 +503,17 @@ def main():
|
||||
}
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
)
|
||||
omni_kwargs = dict(
|
||||
model=args.model,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
@@ -512,15 +524,7 @@ def main():
|
||||
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
|
||||
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
model_class_name=model_class_name,
|
||||
cache_backend=args.cache_backend,
|
||||
|
||||
@@ -15,6 +15,7 @@ import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
|
||||
prepare_longcat_video_avatar_model_for_omni,
|
||||
)
|
||||
@@ -235,6 +236,14 @@ def main():
|
||||
dtype="bfloat16",
|
||||
enforce_eager=args.enforce_eager,
|
||||
additional_config=additional_config,
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
ulysses_degree=1,
|
||||
ring_degree=1,
|
||||
cfg_parallel_size=1,
|
||||
tensor_parallel_size=1,
|
||||
vae_patch_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
),
|
||||
)
|
||||
|
||||
audio_input: str | dict[str, str] | None = args.audio
|
||||
|
||||
@@ -23,6 +23,7 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
@@ -249,6 +250,17 @@ def main():
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
ring_degree=args.ring_degree,
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
)
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
|
||||
omni_kwargs = dict(
|
||||
@@ -259,14 +271,7 @@ def main():
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
ring_degree=args.ring_degree,
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
cache_backend=args.cache_backend,
|
||||
|
||||
@@ -38,6 +38,7 @@ from typing import Any
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -305,13 +306,17 @@ def main():
|
||||
print(f" Seed: {args.seed}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
)
|
||||
|
||||
# Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio
|
||||
# is resolved from its config.
|
||||
omni_kwargs: dict[str, Any] = dict(
|
||||
model=args.model,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
parallel_config=parallel_config,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
|
||||
@@ -13,7 +13,6 @@ import yaml
|
||||
from transformers import AutoTokenizer
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
KEY_CFG,
|
||||
KEY_MAX_DECODE_STEPS,
|
||||
@@ -22,6 +21,7 @@ from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
SAMPLE_RATE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
|
||||
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
|
||||
|
||||
try:
|
||||
from .runner import (
|
||||
@@ -142,8 +142,7 @@ def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
|
||||
return _load_speaker_embedding(args.speaker_embedding)
|
||||
if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
|
||||
return None
|
||||
# Runs before the engine is built, so nothing has pre-fetched campplus.onnx.
|
||||
embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths)
|
||||
embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
|
||||
if not embs:
|
||||
raise RuntimeError("Speaker extraction produced no embeddings")
|
||||
return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
|
||||
|
||||
@@ -10,6 +10,7 @@ from typing import Any
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -461,6 +462,17 @@ def main():
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
# Configure parallel settings
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
)
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
|
||||
omni_kwargs = dict(
|
||||
@@ -469,13 +481,7 @@ def main():
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
model_class_name=model_class_name,
|
||||
cache_backend=args.cache_backend,
|
||||
|
||||
@@ -11,6 +11,7 @@ import numpy as np
|
||||
from PIL import Image
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -196,6 +197,14 @@ def main() -> None:
|
||||
},
|
||||
)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
)
|
||||
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
cache_config = {
|
||||
@@ -213,11 +222,7 @@ def main() -> None:
|
||||
|
||||
omni_kwargs = dict(
|
||||
model=args.model,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
parallel_config=parallel_config,
|
||||
model_type=args.model_type,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
|
||||
@@ -416,10 +416,9 @@ For example, keep the first main block's attention projections in BF16 with:
|
||||
```
|
||||
|
||||
The structured option replaces `--quantization fp8`. Online FP8 can be used
|
||||
with H3 layerwise offload and with both DLO transfer paths. The default
|
||||
AllGather path uses the ordinary loader to finalize FP8 weights and scales
|
||||
before sharding them across ranks. `--dlo-no-use-allgather` instead retains
|
||||
complete rank-local tensors and avoids the synchronized request-wave contract.
|
||||
with H3 layerwise offload and with distributed layerwise offload's full-weight
|
||||
per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
|
||||
supported for runtime-created FP8 weights.
|
||||
|
||||
## AMD ROCm (gfx942 / gfx950)
|
||||
|
||||
@@ -871,9 +870,8 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
|
||||
H3 native `tile` mode only.
|
||||
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
|
||||
pure Ulysses.
|
||||
- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model
|
||||
in host memory on every rank during startup before retaining only each rank's
|
||||
shard. Size startup host memory for that transient peak.
|
||||
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
|
||||
use `--dlo-no-use-allgather` when combining online FP8 with DLO.
|
||||
- TeaCache and Cache-DiT cannot be enabled on the same server.
|
||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
||||
- Video Ref2VA accepts one or more video files, but not an additional standalone
|
||||
|
||||
@@ -1,173 +0,0 @@
|
||||
# MiniMax Music 3
|
||||
|
||||
> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: MiniMax
|
||||
- Model: `MiniMaxAI/MiniMax-Music3`
|
||||
- Task: Text-to-music generation (vocals and instrumental together)
|
||||
- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API
|
||||
- Maintainer: @linyueqian
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six
|
||||
minutes. This is not a TTS model with a musical voice: there is no reference
|
||||
speaker, no `voice`, and no `temperature`. What you control is the lyrics, the
|
||||
caption, the seed, and the length.
|
||||
|
||||
## References
|
||||
|
||||
- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3)
|
||||
- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/)
|
||||
- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU),
|
||||
`vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two)
|
||||
|
||||
## Architecture
|
||||
|
||||
| Component | Spec |
|
||||
|---|---|
|
||||
| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k |
|
||||
| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 |
|
||||
| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder |
|
||||
| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 |
|
||||
| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling |
|
||||
| Solver | Euler, 30 steps, acoustic guidance scale 1.7 |
|
||||
| Window | 200 frames per acoustic chunk, 100-frame hop |
|
||||
| Frame rate | 25 frames per second |
|
||||
| Context length | 10,240 tokens |
|
||||
| Output | 32 kHz stereo WAV |
|
||||
|
||||
Generation is two-stage. The backbone predicts one audio frame per decode step,
|
||||
each frame eight RVQ codebooks deep. Every 200 frames are handed to the
|
||||
flow-matching transformer, which solves a latent that the decoder turns into
|
||||
waveform.
|
||||
|
||||
**Guidance doubles the rows.** Every request decodes twice: one row sees the
|
||||
real prompt, the other sees the same prompt with the caption-and-lyrics span
|
||||
replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole
|
||||
song, so a request costs twice the KV its length suggests. Size the pool for
|
||||
rows, not requests.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x H200 141GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12
|
||||
- CUDA: 12.9
|
||||
- vLLM version: 0.27.0
|
||||
- vLLM-Omni version or commit: branch `feat/minimax-music3`
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve MiniMaxAI/MiniMax-Music3 \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--trust-remote-code --omni
|
||||
```
|
||||
|
||||
The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered
|
||||
automatically from the checkpoint's `model_type=minimax_music3`. It colocates
|
||||
both stages on device 0.
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "MiniMaxAI/MiniMax-Music3",
|
||||
"input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me",
|
||||
"instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.",
|
||||
"seed": 1,
|
||||
"max_new_tokens": 750
|
||||
}' \
|
||||
--output song.wav
|
||||
```
|
||||
|
||||
Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames
|
||||
at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a
|
||||
target: the model ends the song itself when it emits the audio-end token, and a
|
||||
response shorter than the cap is the model finishing rather than a truncation.
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone
|
||||
and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights,
|
||||
which sit outside any KV budget.
|
||||
- **Key flags:** none required. Guidance, the solver step count and the window
|
||||
geometry are fixed by the checkpoint and are not request parameters.
|
||||
- **Known limitations:**
|
||||
- The external API is non-streaming; `stream: true` is rejected.
|
||||
- `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a
|
||||
maximum-length song (9,000 frames) cannot both fit. Long captions shorten
|
||||
the maximum song. This is a property of the checkpoint.
|
||||
- The acoustic stage runs in float32. bfloat16 is accepted but measurably
|
||||
degrades the solver.
|
||||
|
||||
### 2x H200 141GB
|
||||
|
||||
Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device:
|
||||
|
||||
```bash
|
||||
vllm serve MiniMaxAI/MiniMax-Music3 \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--trust-remote-code --omni \
|
||||
--deploy-config minimax_music3_2gpu.yaml
|
||||
```
|
||||
|
||||
Only the placement differs; both layouts run the acoustic stage in float32.
|
||||
|
||||
## Writing the request
|
||||
|
||||
Two fields carry the request, and both are required.
|
||||
|
||||
- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`,
|
||||
`[Outro]`) steer the arrangement and are part of the prompt contract.
|
||||
**Put a tag on its own line.** Normalization keeps only the tags on a line
|
||||
that starts with one and drops the rest of that line, so lyrics written next
|
||||
to a tag are silently lost:
|
||||
|
||||
```text
|
||||
"[Verse]\nWalking down the street" -> [start] [verse] Walking down the street
|
||||
"[Verse] Walking down the street" -> [start] [verse]
|
||||
```
|
||||
|
||||
- `instructions` is the **caption** describing genre, instrumentation, tempo,
|
||||
mood and production. It is the strongest control available: vague captions
|
||||
give generic arrangements. Long structured captions work well, since the
|
||||
model was trained on descriptions covering global attributes, emotional
|
||||
progression, vocal detail and arrangement.
|
||||
|
||||
A request is deterministic in its seed: the same lyrics, caption, seed and
|
||||
length return the same audio. Changing only the seed gives another take of the
|
||||
same song.
|
||||
|
||||
## Request parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|---|---|---|---|
|
||||
| `input` | string | required | Lyrics, non-empty. Tags on their own lines |
|
||||
| `instructions` | string | required | Caption, non-empty |
|
||||
| `seed` | int | `0` | Non-negative. Fixes the output for a given request |
|
||||
| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 |
|
||||
| `response_format` | string | `"wav"` | Output container |
|
||||
|
||||
## Parameters the model rejects
|
||||
|
||||
The request is refused rather than silently ignored, so a mistake is visible
|
||||
immediately.
|
||||
|
||||
| Parameter | Reason |
|
||||
|---|---|
|
||||
| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw |
|
||||
| `voice` | There is no speaker to select; the vocal comes from the caption |
|
||||
| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag |
|
||||
| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" |
|
||||
| `stream: true` | The external API is non-streaming |
|
||||
@@ -1,193 +0,0 @@
|
||||
# SANA-WM
|
||||
|
||||
> Camera-controllable first-frame image-to-video world model.
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Efficient-Large-Model / NVlabs SANA
|
||||
- Model: `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` (standard diffusers layout, converted offline from the NVlabs release; Stage-1 transformer + VAE only)
|
||||
- Task: First-frame image-to-video generation with camera control
|
||||
- Mode: Online serving with the OpenAI-compatible video API
|
||||
- Model weights: about 13 GB for the Stage-1 transformer (10 GB) and VAE (2.3 GB)
|
||||
- Local disk: reserve about 40 GB for the Hugging Face cache and runtime artifacts
|
||||
- Recommended GPU: 24 GB or larger CUDA GPU
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want to serve SANA-WM through `/v1/videos` or
|
||||
`/v1/videos/sync`. The model takes a text prompt, a first-frame image, and
|
||||
either an action DSL string or explicit camera poses. vLLM-Omni serves the
|
||||
SANA-WM Stage-1 DiT, decoded through the SANA VAE. The optional LTX-2 refiner
|
||||
stage is not supported by this integration; it is a planned follow-up.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional>
|
||||
- Video API: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA RTX PRO 6000 Blackwell 96GB
|
||||
|
||||
#### Capacity
|
||||
|
||||
- Model storage: the Stage-1 transformer is about 10 GB and the VAE about
|
||||
2.3 GB. The Gemma text encoder is a separate 4.9 GB repo. The engine prefetches
|
||||
the whole model repo at startup (`allow_patterns=["*"]`), which is why the
|
||||
Stage-1 weights live in their own repo — the two-stage one carries an
|
||||
additional 84 GB `refiner/` that this path never loads.
|
||||
- Text encoder: the pipeline tries `google/gemma-2-2b-it` first, then falls back
|
||||
to the ungated mirror `Efficient-Large-Model/gemma-2-2b-it`. The first repo is
|
||||
gated, so without an accepted licence and `hf auth login` you will see one
|
||||
failed load before the fallback succeeds — that is expected, not an error to
|
||||
chase. Set `VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER` to pin a specific repo or a
|
||||
local path and skip both.
|
||||
- Disk sizing: provision about 40 GB of local disk or Hugging Face cache volume
|
||||
so the model, temporary downloads, and generated artifacts fit without cache
|
||||
eviction.
|
||||
- GPU sizing: the default 1280x704, 161-frame, 60-step serving profile peaks at
|
||||
22.6 GB of device memory and takes about 133 s to generate on one RTX PRO
|
||||
6000 Blackwell. The peak lands in the VAE decode, which is why the pipeline
|
||||
forces VAE tiling on regardless of `vae_use_tiling` — without it the same
|
||||
request costs about 9 GB more, and a 321-frame one OOMs outright. On smaller
|
||||
GPUs, lower `width`, `height`, or `num_frames` before serving production
|
||||
requests.
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch
|
||||
build
|
||||
- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni
|
||||
Python environment
|
||||
- vLLM version: Match the repository requirements for your checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
#### Command
|
||||
|
||||
The repo ships the standard Diffusers layout (`model_index.json` +
|
||||
`transformer/`, `vae/`), and its `model_index.json` names `SanaWmPipeline`, so
|
||||
the pipeline class resolves on its own.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
vllm serve BBBBruce/SANA-WM_bidirectional-stage1-diffusers \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
No deploy config: single-stage diffusion models are deliberately absent from
|
||||
`OMNI_PIPELINES` (`vllm_omni/config/pipeline_registry.py`), so stage resolution
|
||||
falls back to the default stage config and a YAML's stage settings — including
|
||||
`default_sampling_params` — would not be applied. The production generation
|
||||
settings therefore live in the model (`num_inference_steps=60`,
|
||||
`guidance_scale=5.0`); a request that omits a field gets them. The examples
|
||||
below still pass every field explicitly so the numbers are visible.
|
||||
|
||||
If you point this at the older two-stage repo
|
||||
(`BBBBruce/SANA-WM_bidirectional-diffusers`), startup fails with `Model class
|
||||
SanaWmTwoStagesPipeline not found in diffusion model registry`, because that
|
||||
repo's `model_index.json` names a class this build does not register. Add
|
||||
`--model-class-name SanaWmPipeline` to override it.
|
||||
|
||||
#### Verification
|
||||
|
||||
Use a short smoke request first:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8091/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "prompt=A slow forward camera move through a quiet city street." \
|
||||
-F "negative_prompt=blurry, low quality, distorted, watermark" \
|
||||
-F "input_reference=@/path/to/first_frame.png;type=image/png" \
|
||||
-F "width=1280" \
|
||||
-F "height=704" \
|
||||
-F "num_frames=9" \
|
||||
-F "fps=16" \
|
||||
-F "num_inference_steps=2" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42" \
|
||||
--form-string 'extra_params={"sana_wm":{"action":"w-8","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
|
||||
-o sana_wm_smoke.mp4
|
||||
```
|
||||
|
||||
For a production-length request, note that the action durations must sum to
|
||||
`num_frames - 1` — the rollout includes the identity start pose — and a
|
||||
mismatch is rejected rather than padded or truncated:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8091/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "prompt=A slow forward camera move through a quiet city street." \
|
||||
-F "negative_prompt=blurry, low quality, distorted, watermark" \
|
||||
-F "input_reference=@/path/to/first_frame.png;type=image/png" \
|
||||
-F "width=1280" \
|
||||
-F "height=704" \
|
||||
-F "num_frames=161" \
|
||||
-F "fps=16" \
|
||||
-F "num_inference_steps=60" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42" \
|
||||
--form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
|
||||
-o sana_wm_output.mp4
|
||||
```
|
||||
|
||||
Use `POST /v1/videos` instead when you want job storage and polling rather than
|
||||
inline MP4 bytes. It accepts the same form fields as `/v1/videos/sync`.
|
||||
|
||||
```bash
|
||||
create_response=$(curl -sS -X POST http://localhost:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A slow forward camera move through a quiet city street." \
|
||||
-F "negative_prompt=blurry, low quality, distorted, watermark" \
|
||||
-F "input_reference=@/path/to/first_frame.png;type=image/png" \
|
||||
-F "width=1280" \
|
||||
-F "height=704" \
|
||||
-F "num_frames=161" \
|
||||
-F "fps=16" \
|
||||
-F "num_inference_steps=60" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42" \
|
||||
--form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}')
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
curl -sS "http://localhost:8091/v1/videos/${video_id}" | jq .
|
||||
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o sana_wm_output.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Sequence parallelism is not supported. The bidirectional gated delta
|
||||
recurrence carries state across frames, so a rank cannot denoise a slice
|
||||
of the token sequence in isolation; supporting it needs a distributed scan
|
||||
or an all-gather before the GDN blocks.
|
||||
|
||||
- `input_reference` is required for the first frame. Use `image_reference` only
|
||||
when you need a JSON-safe image URL or data URL instead of a multipart upload.
|
||||
- `sana_wm` must provide exactly one of `action` or `camera`.
|
||||
- Action strings use comma-separated `<keys>-<duration>` segments. Supported
|
||||
keys are `w`, `a`, `s`, `d` for translation and `i`, `j`, `k`, `l` for
|
||||
pitch/yaw rotation. The durations must sum to `num_frames - 1`.
|
||||
- Explicit camera control (alternative to `action`): pass
|
||||
`"camera": {"poses": [...]}` where `poses` is a list of `num_frames`
|
||||
camera-to-world 4x4 matrices (row-major, OpenCV `+X right, +Y down, +Z forward`
|
||||
convention), e.g.
|
||||
`extra_params={"sana_wm":{"camera":{"poses":[[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], ...]},"intrinsics":{...}}}`.
|
||||
Most callers should prefer `action`; explicit poses exist for callers that
|
||||
already have a per-frame trajectory.
|
||||
- Explicit `intrinsics` are recommended and take the mapping form
|
||||
`{"fx":640,"fy":640,"cx":640,"cy":352}` (for 1280x704). This `{fx,fy,cx,cy}`
|
||||
mapping is the only accepted intrinsics form; omit `intrinsics` to derive them
|
||||
from the output resolution. All four values must be finite, and `fx`/`fy`
|
||||
must be positive — the ray map divides by them.
|
||||
- The video API returns decoded MP4 bytes and has no `output_type` field, so
|
||||
raw Stage-1 latents are reachable only from the offline API
|
||||
(`OmniDiffusionSamplingParams(output_type="latent")`); see
|
||||
[`tests/e2e/offline_inference/test_sana_wm.py`](../../tests/e2e/offline_inference/test_sana_wm.py).
|
||||
Putting `output_type` in `extra_params` does not work: it lands in
|
||||
`sampling_params.extra_args`, while the pipeline reads the top-level field.
|
||||
@@ -70,7 +70,6 @@ recipes/
|
||||
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
|
||||
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
|
||||
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
|
||||
| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
|
||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||
|
||||
@@ -189,10 +189,8 @@ wait
|
||||
|
||||
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
|
||||
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
|
||||
- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader
|
||||
instead of direct checkpoint mmap, then shard the finalized FP8 weights and
|
||||
scales. Size host memory for a transient complete FP8 model on every rank at
|
||||
startup
|
||||
- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
|
||||
use `--dlo-no-use-allgather` or disable online quantization
|
||||
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
|
||||
`--dlo-no-use-allgather` retains the standard TP loader and should be treated
|
||||
as an experimental compatibility path
|
||||
|
||||
@@ -5,21 +5,20 @@
|
||||
## Summary
|
||||
|
||||
- Vendor: inclusionAI
|
||||
- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B`
|
||||
- Model: `inclusionAI/Ming-omni-tts-0.5B`
|
||||
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
|
||||
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API
|
||||
- Maintainer: Community
|
||||
|
||||
## References
|
||||
|
||||
- [Huggingface: inclusionAI/Ming-omni-tts-16.8B-A3B](https://huggingface.co/inclusionAI/Ming-omni-tts-16.8B-A3B)
|
||||
- [Huggingface: inclusionAI/Ming-omni-tts-0.5B](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
|
||||
- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
|
||||
- Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts)
|
||||
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B and the MoE 16.8B two-stage TTS pipeline deployment.
|
||||
This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
|
||||
Other hardware is welcome as community validation lands.
|
||||
|
||||
## CUDA
|
||||
@@ -40,7 +39,11 @@ Other hardware is welcome as community validation lands.
|
||||
Launch the two-stage talker:
|
||||
|
||||
```bash
|
||||
vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091
|
||||
vllm serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
#### Verification
|
||||
@@ -51,7 +54,7 @@ Basic synthesis (save the WAV bytes):
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "inclusionAI/Ming-omni-tts-16.8B-A3B",
|
||||
"model": "inclusionAI/Ming-omni-tts-0.5B",
|
||||
"input": "你好,这是 Ming 在线语音合成测试。",
|
||||
"response_format": "wav"
|
||||
}' --output ming_tts_basic.wav
|
||||
@@ -71,7 +74,7 @@ Reference-audio zero-shot cloning:
|
||||
REF_AUDIO="$BASE/10002287-00000094.wav"
|
||||
jq -n --arg ref_audio "$REF_AUDIO" \
|
||||
'{
|
||||
model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
|
||||
model: "inclusionAI/Ming-omni-tts-0.5B",
|
||||
input: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。",
|
||||
ref_audio: $ref_audio,
|
||||
ref_text: "在此奉劝大家别乱打美白针。",
|
||||
@@ -104,7 +107,7 @@ jq -n \
|
||||
--arg a "$REF_A" \
|
||||
--arg b "$REF_B" \
|
||||
'{
|
||||
model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
|
||||
model: "inclusionAI/Ming-omni-tts-0.5B",
|
||||
input: $input,
|
||||
ref_audio: [$a, $b],
|
||||
ref_text: $ref_text,
|
||||
|
||||
@@ -9,17 +9,13 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
|
||||
DailyOmniDataset,
|
||||
daily_omni_local_qa_json,
|
||||
daily_omni_local_videos_dir,
|
||||
ensure_daily_omni_hub_root,
|
||||
resolve_daily_omni_local_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
_HUB_ID = "liarliar/Daily-Omni"
|
||||
|
||||
|
||||
def _write_qa_json(root: Path) -> Path:
|
||||
qa = root / "qa.json"
|
||||
@@ -27,28 +23,6 @@ def _write_qa_json(root: Path) -> Path:
|
||||
return qa
|
||||
|
||||
|
||||
def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path:
|
||||
"""Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access.
|
||||
|
||||
Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it
|
||||
only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads.
|
||||
"""
|
||||
constants = pytest.importorskip("huggingface_hub").constants
|
||||
|
||||
cache = tmp_path / "hub"
|
||||
repo_dir = cache / "datasets--liarliar--Daily-Omni"
|
||||
snapshot = repo_dir / "snapshots" / "deadbeef"
|
||||
snapshot.mkdir(parents=True)
|
||||
if with_qa:
|
||||
_write_qa_json(snapshot)
|
||||
(repo_dir / "refs").mkdir()
|
||||
(repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
|
||||
|
||||
monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
|
||||
monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
|
||||
return snapshot
|
||||
|
||||
|
||||
def test_resolve_local_root_returns_none_for_hub_id() -> None:
|
||||
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
|
||||
assert resolve_daily_omni_local_root("") is None
|
||||
@@ -116,37 +90,3 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
|
||||
root = resolve_daily_omni_local_root(str(tmp_path))
|
||||
assert daily_omni_local_qa_json(root) is None
|
||||
assert daily_omni_local_videos_dir(root) is None
|
||||
|
||||
|
||||
def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
|
||||
|
||||
root = ensure_daily_omni_hub_root(_HUB_ID)
|
||||
|
||||
assert root == snapshot.resolve()
|
||||
assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
|
||||
|
||||
|
||||
def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""The regression this file exists for: a Hub id must not require ``load_dataset``."""
|
||||
snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
|
||||
# Any use of `datasets` here means we went back to the path that cannot read the hub cache.
|
||||
monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
|
||||
|
||||
dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0)
|
||||
|
||||
assert dataset.qa_json_path == (snapshot / "qa.json").resolve()
|
||||
assert [row["video_id"] for row in dataset.data] == ["vid0"]
|
||||
|
||||
|
||||
def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False)
|
||||
monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
|
||||
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0)
|
||||
|
||||
message = str(excinfo.value)
|
||||
assert "someone-else/Daily-Omni" in message
|
||||
assert "snapshot_download failed" in message
|
||||
assert "--daily-omni-qa-json" in message
|
||||
@@ -1,119 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Guard the NPU pytest circular import through DiffusionOutput."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def _top_level_import_modules(path: Path) -> set[str]:
|
||||
tree = ast.parse(path.read_text(encoding="utf-8"))
|
||||
modules: set[str] = set()
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.ImportFrom) and node.module:
|
||||
modules.add(node.module)
|
||||
elif isinstance(node, ast.Import):
|
||||
modules.update(alias.name for alias in node.names)
|
||||
return modules
|
||||
|
||||
|
||||
def _assert_isolated_import_succeeds(script: str) -> None:
|
||||
env = os.environ.copy()
|
||||
env.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
env.pop("HIP_VISIBLE_DEVICES", None)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd=REPO_ROOT,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_config_package_does_not_eagerly_import_pipeline_registry() -> None:
|
||||
imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py")
|
||||
assert "vllm_omni.config.config_factory" not in imported
|
||||
assert "vllm_omni.config.pipeline_registry" not in imported
|
||||
|
||||
|
||||
def test_config_package_lazy_exports_still_resolve() -> None:
|
||||
from vllm_omni.config import StageConfigFactory, register_pipeline
|
||||
from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory
|
||||
from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline
|
||||
|
||||
assert StageConfigFactory is DirectStageConfigFactory
|
||||
assert register_pipeline is direct_register_pipeline
|
||||
|
||||
|
||||
def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
|
||||
src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8")
|
||||
init_start = src.index("def __init__(self) -> None:")
|
||||
runtime_start = src.index("def init_diffusion_model_runner_runtime")
|
||||
init_src = src[init_start:runtime_start]
|
||||
runtime_src = src[runtime_start:]
|
||||
assert "apply_minimax_h3_qwen3vl_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
|
||||
|
||||
|
||||
def test_lora_config_import_does_not_load_pipeline_registry() -> None:
|
||||
"""Import LoRAConfig and diffusion.data in a fresh interpreter.
|
||||
|
||||
The NPU failure was ``data.py`` importing LoRAConfig while still loading,
|
||||
then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 →
|
||||
DiffusionOutput. This child process must not preload those exports.
|
||||
"""
|
||||
_assert_isolated_import_succeeds(
|
||||
"""
|
||||
import sys
|
||||
|
||||
from vllm_omni.config.lora import LoRAConfig
|
||||
|
||||
if "vllm_omni.config.pipeline_registry" in sys.modules:
|
||||
raise SystemExit("LoRAConfig import loaded pipeline_registry")
|
||||
if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
|
||||
raise SystemExit("LoRAConfig import loaded PI0_PIPELINE")
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
|
||||
if DiffusionOutput is None or LoRAConfig is None:
|
||||
raise SystemExit("expected LoRAConfig and DiffusionOutput")
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_registry_does_not_import_pi0_runtime() -> None:
|
||||
"""Resolving the pi0 topology must not import its runtime pipeline.
|
||||
|
||||
``pipeline_pi0`` imports ``diffusion.data``. Importing it while
|
||||
``diffusion.data`` is only partially initialized closes the Ascend plugin's
|
||||
circular-import loop, so the registry must depend only on a lightweight
|
||||
pipeline-config module.
|
||||
"""
|
||||
_assert_isolated_import_succeeds(
|
||||
"""
|
||||
import sys
|
||||
|
||||
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
|
||||
|
||||
if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
|
||||
raise SystemExit("pipeline_registry imported the pi0 runtime pipeline")
|
||||
|
||||
pipeline = resolve_pipeline_config("pi0")
|
||||
if pipeline is None or pipeline.model_type != "pi0":
|
||||
raise SystemExit("expected the registered pi0 pipeline config")
|
||||
"""
|
||||
)
|
||||
@@ -5,20 +5,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import fields
|
||||
from inspect import Parameter, signature
|
||||
from pathlib import Path
|
||||
|
||||
import msgspec
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
from pydantic.fields import FieldInfo
|
||||
from transformers import Qwen3OmniMoeConfig
|
||||
from vllm.config import CacheConfig as VllmCacheConfig
|
||||
from vllm.config import CompilationConfig as VllmCompilationConfig
|
||||
from vllm.config import LoadConfig as VllmLoadConfig
|
||||
from vllm.config import ParallelConfig as VllmParallelConfig
|
||||
from vllm.config import ProfilerConfig as VllmProfilerConfig
|
||||
from vllm.config import SchedulerConfig as VllmSchedulerConfig
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config import omni_config as omni_config_module
|
||||
@@ -49,7 +40,6 @@ from vllm_omni.config.stage_config import (
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
|
||||
from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -253,29 +243,9 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
|
||||
],
|
||||
ids=["global", "ar-stage", "generation-stage"],
|
||||
)
|
||||
def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
|
||||
stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
|
||||
|
||||
assert stage.cache_config.cache_dtype == "fp8"
|
||||
assert "cache_dtype" in stage.cache_config._omni_explicit_fields
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
|
||||
def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
|
||||
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
||||
deploy = DeployConfig(
|
||||
stages=[
|
||||
StageDeployConfig(
|
||||
stage_id=0,
|
||||
engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
|
||||
|
||||
assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
|
||||
assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
|
||||
def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
|
||||
with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
|
||||
_from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
|
||||
|
||||
|
||||
def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
|
||||
@@ -467,8 +437,6 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
|
||||
"connector_config",
|
||||
"runtime_config",
|
||||
"parallel_config",
|
||||
"compilation_config",
|
||||
"profiler_config",
|
||||
"quantization_config",
|
||||
}
|
||||
assert "diffusion_config" not in public_fields
|
||||
@@ -487,47 +455,10 @@ def test_runtime_config_fields_match_structured_runtime_scope():
|
||||
"num_gpus",
|
||||
"log_level",
|
||||
"log_stats",
|
||||
"profiler_config",
|
||||
}
|
||||
|
||||
|
||||
def test_upstream_config_compatible_fields_materialize_mapping_inputs():
|
||||
compilation_config = {"backend": "eager"}
|
||||
profiler_config = {"profiler": "cuda"}
|
||||
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
||||
compilation_config=compilation_config,
|
||||
profiler_config=profiler_config,
|
||||
)
|
||||
|
||||
assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
|
||||
assert stage_config.compilation_config.backend == "eager"
|
||||
assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
|
||||
assert stage_config.profiler_config.profiler == "cuda"
|
||||
|
||||
|
||||
def test_upstream_config_compatible_fields_validate_mapping_inputs():
|
||||
with pytest.raises(ValidationError):
|
||||
VllmOmniARStageConfig(
|
||||
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
||||
profiler_config={"delay_iterations": -1},
|
||||
)
|
||||
|
||||
|
||||
def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
|
||||
compilation_config = VllmCompilationConfig(backend="eager")
|
||||
profiler_config = VllmProfilerConfig(profiler="cuda")
|
||||
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
||||
compilation_config=compilation_config,
|
||||
profiler_config=profiler_config,
|
||||
)
|
||||
|
||||
assert stage_config.compilation_config is compilation_config
|
||||
assert stage_config.profiler_config is profiler_config
|
||||
|
||||
|
||||
def test_sub_config_fields_match_structured_scopes():
|
||||
assert {f.name for f in fields(OmniStageModelConfig)} == {
|
||||
"model",
|
||||
@@ -558,6 +489,7 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"enforce_eager",
|
||||
"max_cudagraph_capture_size",
|
||||
"enable_flashinfer_autotune",
|
||||
"compilation_config",
|
||||
"enable_multithread_weight_load",
|
||||
"num_weight_load_threads",
|
||||
"disable_autocast",
|
||||
@@ -567,25 +499,29 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"model_subdir",
|
||||
"tokenizer_subdir",
|
||||
}
|
||||
vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
|
||||
assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
|
||||
assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
|
||||
assert {f.name for f in fields(OmniStageLoadConfig)} == {
|
||||
"tokenizer",
|
||||
"download_dir",
|
||||
"skip_tokenizer_init",
|
||||
"load_format",
|
||||
"tokenizer_mode",
|
||||
"config_format",
|
||||
"skip_mm_profiling",
|
||||
}
|
||||
assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
|
||||
assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
|
||||
assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
|
||||
assert {f.name for f in fields(OmniStageCacheConfig)} == {
|
||||
"kv_cache_memory_bytes",
|
||||
"gpu_memory_utilization",
|
||||
"enable_prefix_caching",
|
||||
"disable_hybrid_kv_cache_manager",
|
||||
"mm_processor_cache_gb",
|
||||
"mamba_ssm_cache_dtype",
|
||||
}
|
||||
assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
|
||||
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
|
||||
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
|
||||
"max_num_seqs",
|
||||
"max_num_batched_tokens",
|
||||
"max_model_len",
|
||||
"enable_chunked_prefill",
|
||||
"async_scheduling",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageConnectorConfig)} == {
|
||||
"async_chunk",
|
||||
@@ -594,10 +530,17 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"output_connectors",
|
||||
"input_connectors",
|
||||
}
|
||||
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
|
||||
assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
|
||||
assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
|
||||
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
|
||||
assert {f.name for f in fields(OmniStageParallelConfig)} == {
|
||||
"pipeline_parallel_size",
|
||||
"data_parallel_size",
|
||||
"tensor_parallel_size",
|
||||
"enable_expert_parallel",
|
||||
"world_size",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
|
||||
"pipeline_parallel_size",
|
||||
"data_parallel_size",
|
||||
"tensor_parallel_size",
|
||||
"sequence_parallel_size",
|
||||
"ulysses_degree",
|
||||
"ring_degree",
|
||||
@@ -611,86 +554,20 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"mask_sp_padding",
|
||||
"hsdp_shard_size",
|
||||
"hsdp_replicate_size",
|
||||
"enable_expert_parallel",
|
||||
"world_size",
|
||||
}
|
||||
|
||||
|
||||
def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
|
||||
scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
|
||||
assert scheduler_config.max_num_encoder_input_tokens == 4096
|
||||
assert scheduler_config.encoder_cache_size == 4096
|
||||
|
||||
unresolved_scheduler_config = OmniStageSchedulerConfig()
|
||||
assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
|
||||
assert unresolved_scheduler_config.encoder_cache_size is None
|
||||
|
||||
parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
|
||||
assert parallel_config.data_parallel_index == 2
|
||||
|
||||
diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
|
||||
assert diffusion_parallel_config.data_parallel_index == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[
|
||||
OmniStageLoadConfig,
|
||||
OmniStageCacheConfig,
|
||||
OmniStageSchedulerConfig,
|
||||
OmniStageParallelConfig,
|
||||
OmniStageDiffusionParallelConfig,
|
||||
],
|
||||
)
|
||||
def test_inherited_sub_configs_are_keyword_only(config_cls):
|
||||
assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
config_cls(1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[
|
||||
OmniStageLoadConfig,
|
||||
OmniStageCacheConfig,
|
||||
OmniStageSchedulerConfig,
|
||||
OmniStageParallelConfig,
|
||||
OmniStageDiffusionParallelConfig,
|
||||
],
|
||||
)
|
||||
def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
|
||||
config = config_cls()
|
||||
|
||||
uninitialized_fields = [
|
||||
config_field.name
|
||||
for config_field in fields(config)
|
||||
if isinstance(getattr(config, config_field.name), FieldInfo)
|
||||
]
|
||||
assert uninitialized_fields == []
|
||||
|
||||
msgspec.msgpack.encode(_serialize_stage_config(config))
|
||||
|
||||
|
||||
def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
|
||||
omni_config = _from_pipeline_key("qwen3_tts")
|
||||
|
||||
for stage_config in omni_config.stage_configs:
|
||||
payload = {
|
||||
"stage_id": stage_config.stage_id,
|
||||
"stage_config": _serialize_stage_config(stage_config),
|
||||
}
|
||||
msgspec.msgpack.encode(payload)
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_fields_cover_legacy_surface():
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
|
||||
structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
|
||||
expected_upstream_fields = {"mask_sp_padding"}
|
||||
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
|
||||
|
||||
assert legacy_fields | expected_upstream_fields <= structured_fields
|
||||
assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
|
||||
assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
|
||||
@@ -733,7 +610,8 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
|
||||
|
||||
def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
|
||||
cfg = OmniStageDiffusionParallelConfig(
|
||||
ulysses_degree=4,
|
||||
pipeline_parallel_size=2,
|
||||
ulysses_degree=2,
|
||||
use_hsdp=True,
|
||||
hsdp_shard_size=-1,
|
||||
hsdp_replicate_size=2,
|
||||
@@ -744,10 +622,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
|
||||
with pytest.raises(ValueError, match="not compatible with TP"):
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
|
||||
|
||||
with pytest.raises(ValueError, match="not compatible with DP"):
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
|
||||
|
||||
|
||||
@@ -1200,21 +1078,3 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
|
||||
|
||||
assert cfg.quantization_config == quantization_config
|
||||
|
||||
|
||||
def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
|
||||
# Exercise the terminal construction performed by the future typed startup
|
||||
# path without probing ports or loading remote model metadata.
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
|
||||
cfg.enrich_config()
|
||||
|
||||
assert cfg.quantization_config is not None
|
||||
assert cfg.quantization_config.get_name() == "int8"
|
||||
@@ -7,10 +7,12 @@ The detailed RFC document is maintained outside this repository (local/internal
|
||||
- Qwen3 tests: `test_reliability_qwen3_omni.py`
|
||||
- Wan2.2 tests: `test_reliability_wan22.py`
|
||||
- HunyuanImage tests: `test_reliability_hunyuan_image.py`
|
||||
- VoxCPM2 tests: `test_reliability_voxcpm2.py`
|
||||
|
||||
```bash
|
||||
pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
```
|
||||
@@ -976,68 +976,3 @@ def test_voices_delete_invalid_requests(
|
||||
"err_message": err_message,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# ─── POST /v1/audio/speech · MiniMax Music 3 ───
|
||||
|
||||
# Text-to-music on the speech endpoint. Most of the speech contract does not
|
||||
# apply: there is no speaker, no reference audio and no temperature, and the
|
||||
# lyrics alone are not a request. Each case below pins a rejection so a caller
|
||||
# who assumes the usual TTS shape finds out immediately.
|
||||
_MINIMAX_MUSIC3_SPEECH = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model="MiniMaxAI/MiniMax-Music3",
|
||||
stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
|
||||
server_args=_SPEECH_SERVER_ARGS,
|
||||
),
|
||||
id="minimax_music3",
|
||||
marks=hardware_marks(res={"cuda": "H100"}),
|
||||
),
|
||||
]
|
||||
|
||||
_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight"
|
||||
_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor."
|
||||
# 25 frames per second, capped at six minutes of song.
|
||||
_MINIMAX_MUSIC3_MAX_FRAMES = 9000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides, err_message",
|
||||
[
|
||||
pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
|
||||
pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"),
|
||||
pytest.param({"input": " "}, ("input",), id="lyrics_blank"),
|
||||
pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
|
||||
pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
|
||||
pytest.param(
|
||||
{"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1},
|
||||
("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)),
|
||||
id="max_new_tokens_over_cap",
|
||||
),
|
||||
pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
|
||||
def test_minimax_music3_speech_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
overrides: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
body: dict[str, Any] = {
|
||||
"model": omni_server.model,
|
||||
"input": _MINIMAX_MUSIC3_LYRICS,
|
||||
"instructions": _MINIMAX_MUSIC3_CAPTION,
|
||||
"seed": 1,
|
||||
"max_new_tokens": 250,
|
||||
"response_format": "wav",
|
||||
}
|
||||
for key, value in overrides.items():
|
||||
if value is None:
|
||||
body.pop(key, None)
|
||||
else:
|
||||
body[key] = value
|
||||
openai_client.send_audio_speech_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
@@ -71,7 +71,7 @@ def _default_oom_device_spec() -> str:
|
||||
|
||||
OOM_INJECTION_CONFIG = {
|
||||
"device": _default_oom_device_spec(),
|
||||
"target_mem_ratio": 0.95,
|
||||
"target_mem_ratio": 0.92,
|
||||
"hold_seconds": 0,
|
||||
"startup_timeout_sec": 20,
|
||||
"strict": False,
|
||||
|
||||
@@ -0,0 +1,508 @@
|
||||
"""
|
||||
VoxCPM2 reliability integration tests.
|
||||
|
||||
Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import concurrent.futures
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.dfx.conftest import (
|
||||
assert_fault_exception,
|
||||
create_reliability_omni_server_params,
|
||||
resolve_oom_device_spec,
|
||||
)
|
||||
from tests.dfx.reliability.helpers import (
|
||||
FaultInjector,
|
||||
assert_no_server_tree_process_residual_and_gpu_release,
|
||||
extract_openai_error_contract_from_bytes,
|
||||
get_health_raw,
|
||||
inject_gpu_oom,
|
||||
make_process_kill_fault_injector,
|
||||
make_server_root_kill_fault_injector,
|
||||
make_server_tree_kill_fault_injector,
|
||||
post_json_raw,
|
||||
run_fault_injection_with_rate_load,
|
||||
stop_gpu_oom_hogs,
|
||||
worker_residual_timeout_after_kill_signal,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OpenAIClientHandler
|
||||
|
||||
RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
|
||||
{
|
||||
"test_name": "voxcpm2_reliability_default",
|
||||
"server_params": {
|
||||
"model": "openbmb/VoxCPM2",
|
||||
"stage_config_name": "voxcpm2.yaml",
|
||||
# VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
|
||||
# would otherwise override deploy YAML settings during server startup.
|
||||
"server_args": ["--trust-remote-code", "--disable-log-stats"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
|
||||
OOM_INJECTION_CONFIG = {
|
||||
"target_mem_ratio": 0.96,
|
||||
"hold_seconds": 0,
|
||||
"startup_timeout_sec": 20,
|
||||
"strict": False,
|
||||
}
|
||||
FAULT_ERROR_KEYWORDS = (
|
||||
"the request failed",
|
||||
"oom",
|
||||
"out of memory",
|
||||
"cuda",
|
||||
"orchestrator",
|
||||
"timeout",
|
||||
"connection",
|
||||
"500",
|
||||
"503",
|
||||
)
|
||||
PROCESS_KILL_ERROR_KEYWORDS = (
|
||||
"timeout",
|
||||
"did not complete within",
|
||||
"connection",
|
||||
"engine",
|
||||
"orchestrator",
|
||||
"dead",
|
||||
"internal",
|
||||
"500",
|
||||
"503",
|
||||
)
|
||||
RUNTIME_WORKER_PATTERN = "VLLM::"
|
||||
SERVE_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGINT", id="sigint"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
SERVE_WITH_LOAD_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGINT", id="sigint"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
TREE_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
TREE_WITH_LOAD_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
WORKER_SIGNAL_FAULT_PARAMS = [
|
||||
pytest.param(
|
||||
make_process_kill_fault_injector(
|
||||
grep_patterns=RUNTIME_WORKER_PATTERN,
|
||||
signal_name="SIGTERM",
|
||||
limit=1,
|
||||
post_kill_wait_seconds=2.0,
|
||||
),
|
||||
id="runtime_process_chain_sigterm",
|
||||
),
|
||||
pytest.param(
|
||||
make_process_kill_fault_injector(
|
||||
grep_patterns=RUNTIME_WORKER_PATTERN,
|
||||
signal_name="SIGKILL",
|
||||
limit=1,
|
||||
post_kill_wait_seconds=2.0,
|
||||
),
|
||||
id="runtime_process_chain_sigkill",
|
||||
),
|
||||
]
|
||||
|
||||
VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
|
||||
INFLIGHT_INJECTION_REQUEST_RATE = 0.3
|
||||
INFLIGHT_INJECTION_REQUEST_COUNT = 10
|
||||
_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
|
||||
_SPEECH_FF_MAX_ELAPSED_SEC = 30
|
||||
|
||||
|
||||
def _speech_request_config(omni_server: Any) -> dict[str, Any]:
|
||||
return {
|
||||
"model": omni_server.model,
|
||||
"input": "The weather is nice today, perfect for a walk in the park.",
|
||||
"stream": False,
|
||||
"timeout": 120.0,
|
||||
"response_format": "wav",
|
||||
"voice": "default",
|
||||
}
|
||||
|
||||
|
||||
def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
|
||||
openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
|
||||
|
||||
|
||||
def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
|
||||
base = "The weather is nice today, perfect for a walk in the park. "
|
||||
return {
|
||||
"model": omni_server.model,
|
||||
"input": base * 80,
|
||||
"max_new_tokens": 4096,
|
||||
"ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
|
||||
"stream": False,
|
||||
"timeout": 300.0,
|
||||
"response_format": "wav",
|
||||
"voice": "default",
|
||||
}
|
||||
|
||||
|
||||
def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
|
||||
openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
|
||||
|
||||
|
||||
def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": "post-fault fast-fail check",
|
||||
"voice": "default",
|
||||
}
|
||||
start = time.monotonic()
|
||||
try:
|
||||
status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
|
||||
assert status >= 500, (
|
||||
f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
|
||||
|
||||
|
||||
def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
|
||||
deadline = time.monotonic() + 20.0
|
||||
last_observation = ""
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
status, body = get_health_raw(host, port, timeout_sec=5)
|
||||
last_observation = f"http={status}, body={body[:200]!r}"
|
||||
if status == 503:
|
||||
return
|
||||
except Exception as exc: # noqa: BLE001
|
||||
last_observation = f"exception={exc!r}"
|
||||
return
|
||||
time.sleep(0.5)
|
||||
pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_speech_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
) -> None:
|
||||
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
|
||||
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
|
||||
handle = inject_gpu_oom(
|
||||
device=device_spec,
|
||||
target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
|
||||
hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
|
||||
startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
|
||||
strict=OOM_INJECTION_CONFIG["strict"],
|
||||
)
|
||||
try:
|
||||
try:
|
||||
_submit_oom_speech_request(openai_client_function, omni_server_function)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
|
||||
else:
|
||||
pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
|
||||
finally:
|
||||
stop_gpu_oom_hogs(handle)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_speech_error_contract(
|
||||
omni_server_function,
|
||||
) -> None:
|
||||
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
|
||||
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
|
||||
handle = inject_gpu_oom(
|
||||
device=device_spec,
|
||||
target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
|
||||
hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
|
||||
startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
|
||||
strict=OOM_INJECTION_CONFIG["strict"],
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
oom_config = _oom_speech_request_config(omni_server_function)
|
||||
payload = {
|
||||
"model": oom_config["model"],
|
||||
"input": oom_config["input"],
|
||||
"voice": oom_config["voice"],
|
||||
"max_new_tokens": oom_config["max_new_tokens"],
|
||||
"ref_audio": oom_config["ref_audio"],
|
||||
}
|
||||
try:
|
||||
status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
|
||||
finally:
|
||||
stop_gpu_oom_hogs(handle)
|
||||
|
||||
assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
|
||||
speech_error = extract_openai_error_contract_from_bytes(body)
|
||||
assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
|
||||
assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize(
|
||||
"fault_injector",
|
||||
WORKER_SIGNAL_FAULT_PARAMS,
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_request_failure(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
) -> None:
|
||||
try:
|
||||
_submit_speech_request(openai_client_function, omni_server_after_fault_function)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
else:
|
||||
pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize(
|
||||
"fault_injector",
|
||||
WORKER_SIGNAL_FAULT_PARAMS,
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_fast_fail_and_concurrent(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
) -> None:
|
||||
host = omni_server_after_fault_function.host
|
||||
port = omni_server_after_fault_function.port
|
||||
model = omni_server_after_fault_function.model
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": "Say hello in one short sentence.",
|
||||
"voice": "default",
|
||||
}
|
||||
|
||||
start = time.monotonic()
|
||||
try:
|
||||
status, body = post_json_raw(
|
||||
host,
|
||||
port,
|
||||
"/v1/audio/speech",
|
||||
payload,
|
||||
timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
|
||||
f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
|
||||
)
|
||||
assert status >= 500, (
|
||||
"[process_kill fast_fail] expected server-side failure after fault, "
|
||||
f"got status={status}, body={body[:200]!r}"
|
||||
)
|
||||
except Exception:
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
|
||||
f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
|
||||
)
|
||||
|
||||
start = time.monotonic()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
|
||||
futures = [
|
||||
executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
|
||||
for _ in range(3)
|
||||
]
|
||||
done, pending = concurrent.futures.wait(
|
||||
futures,
|
||||
timeout=30,
|
||||
return_when=concurrent.futures.ALL_COMPLETED,
|
||||
)
|
||||
|
||||
elapsed = time.monotonic() - start
|
||||
assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
|
||||
assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
|
||||
|
||||
fault_observed = False
|
||||
for future in done:
|
||||
try:
|
||||
future.result()
|
||||
except Exception as exc:
|
||||
fault_observed = True
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
assert fault_observed, (
|
||||
"[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize(
|
||||
"fault_injector",
|
||||
WORKER_SIGNAL_FAULT_PARAMS,
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_worker_with_load_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
fault_injector: FaultInjector,
|
||||
) -> None:
|
||||
scenario = "kill_worker_with_load"
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
|
||||
inject_fault=lambda: fault_injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
completion_timeout_sec=120.0,
|
||||
)
|
||||
assert load_result["failure_observed"], (
|
||||
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
injector(omni_server_function)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
completion_timeout_sec=300.0,
|
||||
)
|
||||
assert load_result["failure_observed"], (
|
||||
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
|
||||
injector = make_server_tree_kill_fault_injector(
|
||||
signal_name=signal_name,
|
||||
post_kill_wait_seconds=2.0,
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
injector(omni_server_function)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
|
||||
injector = make_server_tree_kill_fault_injector(
|
||||
signal_name=signal_name,
|
||||
post_kill_wait_seconds=2.0,
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
completion_timeout_sec=120.0,
|
||||
)
|
||||
assert load_result["failure_observed"], (
|
||||
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
|
||||
|
||||
Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
|
||||
|
||||
STABILITY_DIR = Path(__file__).resolve().parent.parent
|
||||
DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
|
||||
CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
|
||||
DEFAULT_NUM_PROMPTS_PER_BATCH = 20
|
||||
STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
|
||||
|
||||
try:
|
||||
BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
|
||||
def test_stability_voxcpm2(omni_server, stability_benchmark_params):
|
||||
test_name = stability_benchmark_params["test_name"]
|
||||
params = stability_benchmark_params["params"]
|
||||
duration_sec = params.get("duration_sec", 300)
|
||||
num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
|
||||
request_rate = params.get("request_rate")
|
||||
max_concurrency = params.get("max_concurrency")
|
||||
|
||||
bench_params = {
|
||||
k: v
|
||||
for k, v in params.items()
|
||||
if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
|
||||
}
|
||||
|
||||
result = run_stability_benchmark_loop(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
model=omni_server.model,
|
||||
duration_sec=duration_sec,
|
||||
params=bench_params,
|
||||
request_rate=request_rate,
|
||||
max_concurrency=max_concurrency,
|
||||
result_dir=str(STABILITY_DIR),
|
||||
num_prompts_per_batch=num_prompts_per_batch,
|
||||
run_one_batch=_run_one_vllm_bench_batch,
|
||||
)
|
||||
|
||||
assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
|
||||
assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"
|
||||
@@ -1,297 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
import vllm_omni.diffusion.diffusion_kv.initialization as initialization
|
||||
import vllm_omni.diffusion.vllm_config as diffusion_vllm_config
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class _Executor:
|
||||
def __init__(self) -> None:
|
||||
self.specs = [{"layer0": object()}, {"layer0": object()}]
|
||||
self.memory = [1024, 1024]
|
||||
self.configs = None
|
||||
self.profile_requests = None
|
||||
|
||||
def get_kv_cache_specs(self):
|
||||
return self.specs
|
||||
|
||||
def determine_available_kv_memory(self, profile_requests):
|
||||
self.profile_requests = profile_requests
|
||||
return self.memory
|
||||
|
||||
def set_kv_cache_configs(self, configs) -> None:
|
||||
self.configs = configs
|
||||
|
||||
|
||||
def _od_config(**overrides):
|
||||
values = dict(
|
||||
model="test-diffusion",
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
dtype=torch.bfloat16,
|
||||
max_model_len=64,
|
||||
quantization_config=None,
|
||||
tf_model_config=None,
|
||||
enforce_eager=True,
|
||||
is_moe=False,
|
||||
additional_config={},
|
||||
profiler_config=None,
|
||||
parallel_config=SimpleNamespace(
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
data_parallel_size=1,
|
||||
enable_expert_parallel=False,
|
||||
),
|
||||
gpu_memory_utilization=0.9,
|
||||
kv_cache_memory_bytes=None,
|
||||
max_num_batched_tokens=64,
|
||||
num_gpus=1,
|
||||
)
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
|
||||
|
||||
def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
|
||||
executor = _Executor()
|
||||
od_config = _od_config(num_gpus=2)
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=4096),
|
||||
max_in_flight_tokens=128,
|
||||
)
|
||||
worker_configs = [object(), object()]
|
||||
scheduler_kv_cache_config = object()
|
||||
profile_requests = [object()]
|
||||
monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config)
|
||||
build = Mock(return_value=(worker_configs, scheduler_kv_cache_config))
|
||||
monkeypatch.setattr(
|
||||
initialization,
|
||||
"build_native_kv_cache_configs",
|
||||
build,
|
||||
)
|
||||
monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16))
|
||||
|
||||
result = initialization.initialize_diffusion_kv_control_plane(
|
||||
executor,
|
||||
od_config,
|
||||
profile_requests=profile_requests,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
|
||||
assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
|
||||
assert executor.configs == worker_configs
|
||||
assert executor.profile_requests is profile_requests
|
||||
|
||||
|
||||
def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None:
|
||||
executor = _Executor()
|
||||
od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
|
||||
get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device"))
|
||||
monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device)
|
||||
|
||||
result = initialization.initialize_diffusion_kv_control_plane(executor, od_config)
|
||||
|
||||
assert result is None
|
||||
assert executor.configs is None
|
||||
get_device.assert_not_called()
|
||||
|
||||
|
||||
def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None:
|
||||
executor = _Executor()
|
||||
executor.memory = [1024]
|
||||
od_config = _od_config(num_gpus=2)
|
||||
monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object())
|
||||
|
||||
with pytest.raises(ValueError, match="rank count mismatch"):
|
||||
initialization.initialize_diffusion_kv_control_plane(
|
||||
executor,
|
||||
od_config,
|
||||
profile_requests=[object()],
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_paged_control_plane_requires_prepared_profile_request() -> None:
|
||||
with pytest.raises(ValueError, match="prepared memory profile requests"):
|
||||
initialization.initialize_diffusion_kv_control_plane(
|
||||
_Executor(),
|
||||
_od_config(),
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None:
|
||||
od_config = _od_config()
|
||||
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
num_kv_heads=2,
|
||||
head_size=8,
|
||||
dtype=torch.bfloat16,
|
||||
non_causal=True,
|
||||
)
|
||||
num_blocks = 32
|
||||
|
||||
worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
|
||||
vllm_config,
|
||||
[{"layer0": spec}],
|
||||
[spec.page_size_bytes * num_blocks],
|
||||
)
|
||||
|
||||
assert len(worker_configs) == 1
|
||||
assert worker_configs[0].num_blocks == num_blocks
|
||||
assert scheduler_kv_cache_config.num_blocks == num_blocks
|
||||
assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups
|
||||
|
||||
|
||||
def test_native_control_plane_config_initializes_scheduler() -> None:
|
||||
od_config = _od_config()
|
||||
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
num_kv_heads=2,
|
||||
head_size=8,
|
||||
dtype=torch.bfloat16,
|
||||
non_causal=True,
|
||||
)
|
||||
_, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
|
||||
vllm_config,
|
||||
[{"layer0": spec}],
|
||||
[spec.page_size_bytes * 32],
|
||||
)
|
||||
|
||||
scheduler = RequestScheduler()
|
||||
scheduler.initialize(
|
||||
od_config,
|
||||
kv_cache_config=scheduler_kv_cache_config,
|
||||
scheduler_block_size=spec.block_size,
|
||||
hash_block_size=spec.block_size,
|
||||
kv_vllm_config=vllm_config,
|
||||
)
|
||||
|
||||
assert scheduler._diffusion_kv_manager is not None
|
||||
scheduler.close()
|
||||
|
||||
|
||||
def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None:
|
||||
od_config = _od_config(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
parallel_config=SimpleNamespace(
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=4,
|
||||
data_parallel_size=1,
|
||||
enable_expert_parallel=False,
|
||||
),
|
||||
gpu_memory_utilization=0.5,
|
||||
kv_cache_memory_bytes=4096,
|
||||
max_num_batched_tokens=64,
|
||||
)
|
||||
vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config)
|
||||
original_pp_size = vllm_config.parallel_config.pipeline_parallel_size
|
||||
original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization
|
||||
original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes
|
||||
original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
|
||||
diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config)
|
||||
|
||||
assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size
|
||||
assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization
|
||||
assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes
|
||||
assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget
|
||||
|
||||
|
||||
def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None:
|
||||
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
|
||||
torch.device("cpu"),
|
||||
_od_config(gpu_memory_utilization=0.42),
|
||||
)
|
||||
|
||||
assert vllm_config.cache_config.gpu_memory_utilization == 0.42
|
||||
|
||||
|
||||
def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
|
||||
quantization_config = build_quant_config(
|
||||
{
|
||||
"quant_method": "modelopt",
|
||||
"quant_algo": "FP8",
|
||||
"ignore": [],
|
||||
}
|
||||
)
|
||||
model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
|
||||
_od_config(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
)
|
||||
|
||||
assert model_config.dtype is torch.bfloat16
|
||||
assert model_config.quantization == "modelopt"
|
||||
assert model_config.quantization_config is quantization_config
|
||||
assert model_config.is_quantized
|
||||
assert model_config.is_moe is False
|
||||
assert model_config.is_nvfp4_quantized() is False
|
||||
assert model_config.original_max_model_len == model_config.max_model_len
|
||||
|
||||
unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
|
||||
_od_config(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
quantization_config=None,
|
||||
)
|
||||
)
|
||||
assert not unquantized_model_config.is_quantized
|
||||
|
||||
|
||||
def test_minus_one_max_model_len_uses_model_limit() -> None:
|
||||
od_config = SimpleNamespace(
|
||||
max_model_len=-1,
|
||||
tf_model_config=SimpleNamespace(max_position_embeddings=4096),
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
)
|
||||
|
||||
assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096
|
||||
|
||||
model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
|
||||
_od_config(
|
||||
max_model_len=-1,
|
||||
tf_model_config=SimpleNamespace(max_position_embeddings=4096),
|
||||
)
|
||||
)
|
||||
assert model_config.max_model_len == 4096
|
||||
assert model_config.original_max_model_len == -1
|
||||
|
||||
|
||||
def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None:
|
||||
vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
|
||||
torch.device("cpu"),
|
||||
_od_config(
|
||||
max_model_len=-1,
|
||||
tf_model_config=SimpleNamespace(max_position_embeddings=4096),
|
||||
),
|
||||
)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
num_kv_heads=2,
|
||||
head_size=8,
|
||||
dtype=torch.bfloat16,
|
||||
non_causal=True,
|
||||
)
|
||||
|
||||
initialization.build_native_kv_cache_configs(
|
||||
vllm_config,
|
||||
[{"layer0": spec}],
|
||||
[spec.page_size_bytes * 16],
|
||||
)
|
||||
|
||||
assert vllm_config.model_config.original_max_model_len == -1
|
||||
assert vllm_config.model_config.max_model_len == 256
|
||||
@@ -1,211 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
|
||||
|
||||
from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
|
||||
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
BLOCK_SIZE = 4
|
||||
|
||||
|
||||
def _config(num_blocks: int) -> KVCacheConfig:
|
||||
native_kv_managers.register_all_kvcache_specs(None)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=BLOCK_SIZE,
|
||||
num_kv_heads=2,
|
||||
head_size=8,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
return KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
|
||||
kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
|
||||
)
|
||||
|
||||
|
||||
def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
|
||||
return DiffusionKVRequest(
|
||||
f"{public_id}/diffusion-kv/{sequence_id}",
|
||||
sequence_id=sequence_id,
|
||||
prefix_len=4,
|
||||
target_len=4,
|
||||
seq_len=seq_len,
|
||||
kv_contexts=kv_contexts,
|
||||
)
|
||||
|
||||
|
||||
def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
|
||||
return DiffusionKVCacheManager(
|
||||
_config(num_blocks),
|
||||
max_model_len=max_model_len,
|
||||
scheduler_block_size=BLOCK_SIZE,
|
||||
hash_block_size=BLOCK_SIZE,
|
||||
)
|
||||
|
||||
|
||||
def test_reserve_and_free_multi_cfg_request() -> None:
|
||||
manager = _manager(8)
|
||||
free_before = manager.native_manager.block_pool.get_num_free_blocks()
|
||||
requests = (_request("public", 0), _request("public", 1))
|
||||
|
||||
metadata = manager.reserve_request("public", requests)
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.request_id == "public"
|
||||
assert metadata.allocation_generation == 1
|
||||
assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1]
|
||||
assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4]
|
||||
assert [sequence.target_len for sequence in metadata.sequences] == [4, 4]
|
||||
assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8]
|
||||
assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
|
||||
assert all(sequence.context_ids == () for sequence in metadata.sequences)
|
||||
assert metadata.contexts == ()
|
||||
assert metadata == manager.get_metadata("public")
|
||||
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
|
||||
|
||||
manager.free_request("public")
|
||||
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
|
||||
|
||||
|
||||
def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None:
|
||||
manager = _manager(3)
|
||||
free_before = manager.native_manager.block_pool.get_num_free_blocks()
|
||||
|
||||
with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
|
||||
manager.reserve_request(
|
||||
"public",
|
||||
(_request("public", 0), _request("public", 1)),
|
||||
)
|
||||
|
||||
assert manager.has_request("public") is False
|
||||
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
|
||||
|
||||
|
||||
def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None:
|
||||
manager = _manager(5)
|
||||
assert manager.reserve_request("running", (_request("running", 0),)) is not None
|
||||
free_before = manager.native_manager.block_pool.get_num_free_blocks()
|
||||
|
||||
with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"):
|
||||
manager.reserve_request(
|
||||
"impossible",
|
||||
tuple(_request("impossible", sequence_id) for sequence_id in range(3)),
|
||||
)
|
||||
|
||||
assert manager.has_request("impossible") is False
|
||||
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
|
||||
|
||||
|
||||
def test_temporary_capacity_pressure_returns_none() -> None:
|
||||
manager = _manager(3)
|
||||
assert manager.reserve_request("running", (_request("running", 0),)) is not None
|
||||
|
||||
allocation = manager.reserve_request("waiting", (_request("waiting", 0),))
|
||||
|
||||
assert allocation is None
|
||||
assert manager.has_request("waiting") is False
|
||||
|
||||
|
||||
def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None:
|
||||
manager = _manager(5)
|
||||
assert manager.reserve_request("running", (_request("running", 0),)) is not None
|
||||
waiting = (_request("waiting", 0), _request("waiting", 1))
|
||||
|
||||
assert manager.reserve_request("waiting", waiting) is None
|
||||
assert manager.has_request("waiting") is False
|
||||
|
||||
manager.free_request("running")
|
||||
metadata = manager.reserve_request("waiting", waiting)
|
||||
|
||||
assert metadata is not None
|
||||
assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
|
||||
|
||||
|
||||
def test_reallocated_request_gets_a_new_allocation_generation() -> None:
|
||||
manager = _manager(3)
|
||||
request = (_request("public", 0),)
|
||||
|
||||
first = manager.reserve_request("public", request)
|
||||
assert first is not None
|
||||
manager.free_request("public")
|
||||
second = manager.reserve_request("public", request)
|
||||
|
||||
assert second is not None
|
||||
assert second.allocation_generation == first.allocation_generation + 1
|
||||
|
||||
|
||||
def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None:
|
||||
manager = _manager(8)
|
||||
free_before = manager.native_manager.block_pool.get_num_free_blocks()
|
||||
native_allocate = manager.native_manager.allocate_slots
|
||||
num_calls = 0
|
||||
|
||||
def fail_second_allocation(*args, **kwargs):
|
||||
nonlocal num_calls
|
||||
num_calls += 1
|
||||
if num_calls == 2:
|
||||
raise RuntimeError("injected allocation failure")
|
||||
return native_allocate(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation)
|
||||
|
||||
with pytest.raises(RuntimeError, match="injected allocation failure"):
|
||||
manager.reserve_request(
|
||||
"public",
|
||||
(_request("public", 0), _request("public", 1)),
|
||||
)
|
||||
|
||||
assert manager.has_request("public") is False
|
||||
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
|
||||
|
||||
|
||||
def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
|
||||
manager = _manager(8)
|
||||
context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
|
||||
|
||||
with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
|
||||
manager.reserve_request(
|
||||
"public",
|
||||
(_request("public", 0, kv_contexts=(context,)),),
|
||||
)
|
||||
|
||||
|
||||
def test_rejects_sequence_longer_than_admission_bound() -> None:
|
||||
manager = _manager(8, max_model_len=8)
|
||||
|
||||
with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"):
|
||||
manager.reserve_request("public", (_request("public", 0, seq_len=9),))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"requests",
|
||||
[
|
||||
(_request("public", 1),),
|
||||
(_request("public", 1), _request("public", 0)),
|
||||
(_request("public", 0), _request("public", 2)),
|
||||
],
|
||||
)
|
||||
def test_rejects_ambiguous_sequence_order(requests) -> None:
|
||||
manager = _manager(8)
|
||||
|
||||
with pytest.raises(ValueError, match="ordered by contiguous sequence_id"):
|
||||
manager.reserve_request("public", requests)
|
||||
|
||||
|
||||
def test_close_releases_every_public_request() -> None:
|
||||
manager = _manager(8)
|
||||
free_before = manager.native_manager.block_pool.get_num_free_blocks()
|
||||
assert manager.reserve_request("first", (_request("first", 0),)) is not None
|
||||
assert manager.reserve_request("second", (_request("second", 0),)) is not None
|
||||
|
||||
manager.close()
|
||||
|
||||
assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
|
||||
assert manager.has_request("first") is False
|
||||
assert manager.has_request("second") is False
|
||||
@@ -107,8 +107,11 @@ def test_request_rejects_non_context_values() -> None:
|
||||
def _manager(*, enable_caching: bool) -> KVCacheManager:
|
||||
# vLLM 0.26 lazily registers its built-in KV specs. Another test may have
|
||||
# already registered an out-of-tree spec, making the registry non-empty
|
||||
# before FullAttentionSpec is registered. Initialize the native mappings.
|
||||
native_kv_managers.register_all_kvcache_specs(None)
|
||||
# before FullAttentionSpec is registered. Initialize the native mappings
|
||||
# explicitly; vLLM 0.25 has no registration hook and needs no setup.
|
||||
register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
|
||||
if register_all is not None:
|
||||
register_all(None)
|
||||
|
||||
spec = FullAttentionSpec(
|
||||
block_size=BLOCK_SIZE,
|
||||
|
||||
@@ -1,259 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
|
||||
|
||||
import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class _Backend:
|
||||
@classmethod
|
||||
def indexes_kv_by_block_stride(cls) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def _attention(*, enabled: bool) -> Attention:
|
||||
attention = Attention.__new__(Attention)
|
||||
nn.Module.__init__(attention)
|
||||
attention.paged_kv_cache_role = "primary" if enabled else None
|
||||
attention.paged_kv_cache_dtype = torch.bfloat16
|
||||
attention.num_kv_heads = 2
|
||||
attention.head_size = 8
|
||||
attention.causal = False
|
||||
attention.attn_backend = _Backend
|
||||
return attention
|
||||
|
||||
|
||||
def _runner(attention: Attention) -> DiffusionModelRunner:
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner.vllm_config = SimpleNamespace(
|
||||
cache_config=SimpleNamespace(block_size=4),
|
||||
model_config=SimpleNamespace(dtype=torch.float16),
|
||||
)
|
||||
pipeline = nn.Module()
|
||||
pipeline.image_attention = attention
|
||||
runner.pipeline = pipeline
|
||||
return runner
|
||||
|
||||
|
||||
def test_runner_discovers_native_spec_from_loaded_attention() -> None:
|
||||
specs = _runner(_attention(enabled=True)).get_kv_cache_spec()
|
||||
|
||||
assert set(specs) == {"image_attention"}
|
||||
spec = specs["image_attention"]
|
||||
assert isinstance(spec, FullAttentionSpec)
|
||||
assert spec.block_size == 4
|
||||
assert spec.num_kv_heads == 2
|
||||
assert spec.head_size == 8
|
||||
assert spec.dtype is torch.bfloat16
|
||||
assert spec.indexes_kv_by_block_stride is True
|
||||
assert spec.non_causal is True
|
||||
|
||||
|
||||
def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
|
||||
runner = _runner(_attention(enabled=False))
|
||||
|
||||
with pytest.raises(RuntimeError, match="no cache-enabled Attention"):
|
||||
runner.get_kv_cache_spec()
|
||||
|
||||
|
||||
def test_runner_retains_matching_rank_local_config() -> None:
|
||||
runner = _runner(_attention(enabled=True))
|
||||
spec = runner.get_kv_cache_spec()["image_attention"]
|
||||
config = KVCacheConfig(
|
||||
num_blocks=8,
|
||||
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
|
||||
kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
|
||||
)
|
||||
|
||||
runner.set_kv_cache_config(config)
|
||||
|
||||
assert runner.kv_cache_config is config
|
||||
|
||||
|
||||
def test_runner_rejects_rank_local_config_for_different_layers() -> None:
|
||||
runner = _runner(_attention(enabled=True))
|
||||
spec = runner.get_kv_cache_spec()["image_attention"]
|
||||
config = KVCacheConfig(
|
||||
num_blocks=8,
|
||||
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
|
||||
kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="layer mismatch"):
|
||||
runner.set_kv_cache_config(config)
|
||||
|
||||
|
||||
def test_worker_selects_its_rank_local_config() -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.rank = 1
|
||||
worker.od_config = SimpleNamespace(num_gpus=2)
|
||||
worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
|
||||
configs = [object(), object()]
|
||||
|
||||
worker.set_kv_cache_configs(configs)
|
||||
|
||||
assert worker.installed is configs[1]
|
||||
with pytest.raises(ValueError, match="rank count mismatch"):
|
||||
worker.set_kv_cache_configs(configs[:1])
|
||||
|
||||
|
||||
def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.rank = 0
|
||||
worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096))
|
||||
worker.init_snapshot = SimpleNamespace(free_memory=16384)
|
||||
worker.requested_memory = 8192
|
||||
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024)
|
||||
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
|
||||
profile_request = object()
|
||||
|
||||
assert worker.determine_available_kv_memory(profile_request) == [4096]
|
||||
worker.model_runner.profile_run.assert_called_once_with(profile_request)
|
||||
|
||||
|
||||
def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.rank = 0
|
||||
worker.init_snapshot = SimpleNamespace(free_memory=900)
|
||||
worker.requested_memory = 750
|
||||
worker.vllm_config = SimpleNamespace(
|
||||
cache_config=SimpleNamespace(
|
||||
kv_cache_memory_bytes=0,
|
||||
gpu_memory_utilization=0.75,
|
||||
)
|
||||
)
|
||||
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
|
||||
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
|
||||
profile_result = SimpleNamespace(
|
||||
non_kv_cache_memory=400,
|
||||
after_profile=SimpleNamespace(free_memory=500),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
diffusion_worker_module,
|
||||
"memory_profiling",
|
||||
lambda snapshot, weights_memory: nullcontext(profile_result),
|
||||
)
|
||||
profile_request = object()
|
||||
|
||||
assert worker.determine_available_kv_memory(profile_request) == [350]
|
||||
worker.model_runner.profile_run.assert_called_once_with(profile_request)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
|
||||
def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.rank = 0
|
||||
worker.init_snapshot = SimpleNamespace(free_memory=900)
|
||||
worker.requested_memory = 750
|
||||
worker.vllm_config = SimpleNamespace(
|
||||
cache_config=SimpleNamespace(
|
||||
kv_cache_memory_bytes=None,
|
||||
gpu_memory_utilization=0.75,
|
||||
)
|
||||
)
|
||||
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
|
||||
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
|
||||
profile_result = SimpleNamespace(
|
||||
non_kv_cache_memory=non_kv_cache_memory,
|
||||
after_profile=SimpleNamespace(free_memory=100),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
diffusion_worker_module,
|
||||
"memory_profiling",
|
||||
lambda snapshot, weights_memory: nullcontext(profile_result),
|
||||
)
|
||||
profile_request = object()
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match=(
|
||||
"rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: "
|
||||
f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes"
|
||||
),
|
||||
):
|
||||
worker.determine_available_kv_memory(profile_request)
|
||||
|
||||
worker.model_runner.profile_run.assert_called_once_with(profile_request)
|
||||
|
||||
|
||||
def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.rank = 0
|
||||
worker.init_snapshot = SimpleNamespace(free_memory=1200)
|
||||
worker.requested_memory = 1000
|
||||
worker.vllm_config = SimpleNamespace(
|
||||
cache_config=SimpleNamespace(
|
||||
kv_cache_memory_bytes=None,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
)
|
||||
worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250)
|
||||
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
|
||||
profile_result = SimpleNamespace(
|
||||
non_kv_cache_memory=650,
|
||||
after_profile=SimpleNamespace(free_memory=550),
|
||||
)
|
||||
profiling_calls = []
|
||||
|
||||
def profile(snapshot, *, weights_memory):
|
||||
profiling_calls.append((snapshot, weights_memory))
|
||||
return nullcontext(profile_result)
|
||||
|
||||
monkeypatch.setattr(
|
||||
diffusion_worker_module,
|
||||
"memory_profiling",
|
||||
profile,
|
||||
)
|
||||
profile_request = object()
|
||||
|
||||
assert worker.determine_available_kv_memory(profile_request) == [350]
|
||||
assert profiling_calls == [(worker.init_snapshot, 250)]
|
||||
worker.model_runner.profile_run.assert_called_once_with(profile_request)
|
||||
|
||||
|
||||
def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.init_snapshot = None
|
||||
worker.requested_memory = None
|
||||
worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None))
|
||||
worker.model_runner = SimpleNamespace()
|
||||
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
|
||||
|
||||
with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"):
|
||||
worker.determine_available_kv_memory(object())
|
||||
|
||||
|
||||
def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
|
||||
gathered = []
|
||||
|
||||
def gather(local_result):
|
||||
gathered.append(local_result)
|
||||
return [local_result, (True, "peer-result")]
|
||||
|
||||
def fail_local_probe():
|
||||
raise ValueError("local failure")
|
||||
|
||||
monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather)
|
||||
|
||||
with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"):
|
||||
diffusion_worker_module._run_and_gather_rank_values(
|
||||
"test probe",
|
||||
fail_local_probe,
|
||||
)
|
||||
|
||||
assert gathered == [(False, "ValueError: local failure")]
|
||||
@@ -85,15 +85,19 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
|
||||
assert new_req.diffusion_kv_metadata is metadata
|
||||
|
||||
|
||||
def test_paged_request_without_metadata_fails_before_request_forward() -> None:
|
||||
def test_paged_request_without_metadata_reaches_request_forward() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner._execute_request_list = Mock()
|
||||
expected_output = DiffusionOutput(output=None)
|
||||
runner._execute_request_list = Mock(
|
||||
return_value=SimpleNamespace(
|
||||
runner_outputs=[SimpleNamespace(result=expected_output)],
|
||||
)
|
||||
)
|
||||
req = SimpleNamespace(request_id="req-0")
|
||||
|
||||
with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
|
||||
runner.execute_model(req)
|
||||
assert runner.execute_model(req) is expected_output
|
||||
|
||||
runner._execute_request_list.assert_not_called()
|
||||
runner._execute_request_list.assert_called_once()
|
||||
|
||||
|
||||
def test_dense_request_rejects_metadata_before_request_forward() -> None:
|
||||
@@ -137,18 +141,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_batch_path_rejects_missing_metadata_before_forward() -> None:
|
||||
def test_batch_path_allows_missing_metadata_before_forward() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner._execute_request_list = Mock()
|
||||
expected_output = object()
|
||||
runner._execute_request_list = Mock(return_value=expected_output)
|
||||
new_req = NewRequestData(
|
||||
request_id="req-0",
|
||||
req=SimpleNamespace(request_id="req-0"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
|
||||
runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
|
||||
assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
|
||||
|
||||
runner._execute_request_list.assert_not_called()
|
||||
runner._execute_request_list.assert_called_once()
|
||||
|
||||
|
||||
def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
|
||||
@@ -166,7 +170,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
|
||||
runner._execute_request_list.assert_not_called()
|
||||
|
||||
|
||||
def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
|
||||
def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner.pipeline = object()
|
||||
runner._supports_step_mode = Mock(return_value=False)
|
||||
@@ -175,10 +179,10 @@ def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
|
||||
req=SimpleNamespace(request_id="req-0"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
|
||||
with pytest.raises(ValueError, match="does not support step execution"):
|
||||
runner.execute_stepwise(make_scheduler_output(new_req))
|
||||
|
||||
runner._supports_step_mode.assert_not_called()
|
||||
runner._supports_step_mode.assert_called_once()
|
||||
|
||||
|
||||
def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
|
||||
@@ -284,42 +288,3 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
|
||||
worker.execute_model(req, SimpleNamespace())
|
||||
|
||||
assert calls == [(req, None)]
|
||||
|
||||
|
||||
def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None:
|
||||
calls: list[tuple] = []
|
||||
|
||||
class ModelRunner:
|
||||
def execute_model(self, req, **kwargs):
|
||||
calls.append((req, kwargs))
|
||||
return DiffusionOutput(output=None)
|
||||
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.model_runner = ModelRunner()
|
||||
worker.lora_manager = None
|
||||
worker._get_profiler = lambda: None
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank",
|
||||
lambda: 1,
|
||||
)
|
||||
req_0 = SimpleNamespace(request_id="req-0")
|
||||
req_1 = SimpleNamespace(request_id="req-1")
|
||||
metadata_0 = make_metadata("req-0")
|
||||
metadata_1 = make_metadata("req-1")
|
||||
envelopes = [
|
||||
NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
|
||||
NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
|
||||
]
|
||||
|
||||
result = worker.execute_model(envelopes, SimpleNamespace())
|
||||
|
||||
assert calls == [
|
||||
(
|
||||
req_1,
|
||||
{
|
||||
"kv_prefetch_job": None,
|
||||
"diffusion_kv_metadata": metadata_1,
|
||||
},
|
||||
)
|
||||
]
|
||||
assert result["dp_rank"] == 1
|
||||
@@ -74,6 +74,14 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
|
||||
lambda: None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])
|
||||
|
||||
@@ -92,8 +92,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
|
||||
monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
|
||||
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
|
||||
monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
|
||||
monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
|
||||
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
@@ -228,8 +229,9 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
|
||||
monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
|
||||
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
|
||||
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
|
||||
monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
|
||||
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
@@ -257,33 +259,6 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch):
|
||||
"""A destroyed diffusion PP group must not block the next initialization."""
|
||||
|
||||
class _DestroyableGroup:
|
||||
def __init__(self) -> None:
|
||||
self.destroy_calls = 0
|
||||
|
||||
def destroy(self) -> None:
|
||||
self.destroy_calls += 1
|
||||
|
||||
pipeline_group = _DestroyableGroup()
|
||||
for name in ("_DP", "_CFG", "_SP"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
monkeypatch.setattr(parallel_state, "_PP", pipeline_group)
|
||||
monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None)
|
||||
for name in ("_DP", "_PCP", "_TP", "_EP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False)
|
||||
|
||||
parallel_state.destroy_model_parallel()
|
||||
|
||||
assert pipeline_group.destroy_calls == 1
|
||||
assert parallel_state.vllm_parallel_state._PP is None
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.core_model
|
||||
def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
|
||||
@@ -320,8 +295,9 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
|
||||
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
|
||||
monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
|
||||
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
|
||||
monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
|
||||
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
|
||||
monkeypatch.setattr(parallel_state, name, None)
|
||||
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
|
||||
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
|
||||
|
||||
@@ -234,7 +234,7 @@ class TestDiffusionParallelConfigHSDP:
|
||||
|
||||
def test_hsdp_cannot_use_with_tp(self):
|
||||
"""Test that HSDP and Tensor Parallelism cannot be used together."""
|
||||
with pytest.raises(ValueError, match="not compatible with TP"):
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
DiffusionParallelConfig(
|
||||
tensor_parallel_size=2,
|
||||
use_hsdp=True,
|
||||
@@ -243,7 +243,7 @@ class TestDiffusionParallelConfigHSDP:
|
||||
|
||||
def test_hsdp_cannot_use_with_dp(self):
|
||||
"""Test that HSDP and Data Parallelism cannot be used together."""
|
||||
with pytest.raises(ValueError, match="not compatible with DP"):
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
DiffusionParallelConfig(
|
||||
data_parallel_size=2,
|
||||
use_hsdp=True,
|
||||
@@ -264,6 +264,158 @@ class TestDiffusionParallelConfigHSDP:
|
||||
assert config.hsdp_shard_size == 2 # auto: 4 // 2
|
||||
|
||||
|
||||
class TestStandaloneHSDPDetection:
|
||||
"""Tests for standalone HSDP detection and dit_parallel_size calculation.
|
||||
|
||||
These tests verify the logic used in initialize_model_parallel() to detect
|
||||
standalone HSDP mode and calculate effective parallel sizes.
|
||||
|
||||
Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def compute_standalone_hsdp_params(
|
||||
data_parallel_size: int = 1,
|
||||
cfg_parallel_size: int = 1,
|
||||
sequence_parallel_size: int = 1,
|
||||
pipeline_parallel_size: int = 1,
|
||||
tensor_parallel_size: int = 1,
|
||||
fully_shard_degree: int = 1,
|
||||
hsdp_replicate_size: int = 1,
|
||||
) -> dict:
|
||||
"""Compute standalone HSDP detection parameters.
|
||||
|
||||
This mirrors the logic in initialize_model_parallel().
|
||||
"""
|
||||
dit_parallel_size = (
|
||||
data_parallel_size
|
||||
* cfg_parallel_size
|
||||
* sequence_parallel_size
|
||||
* pipeline_parallel_size
|
||||
* tensor_parallel_size
|
||||
)
|
||||
|
||||
# Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
|
||||
is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
|
||||
|
||||
# For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
|
||||
if is_standalone_hsdp:
|
||||
effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
|
||||
else:
|
||||
effective_dit_parallel_size = dit_parallel_size
|
||||
|
||||
effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
|
||||
|
||||
return {
|
||||
"original_dit_parallel_size": dit_parallel_size,
|
||||
"is_standalone_hsdp": is_standalone_hsdp,
|
||||
"effective_dit_parallel_size": effective_dit_parallel_size,
|
||||
"effective_dp_size": effective_dp_size,
|
||||
}
|
||||
|
||||
def test_standalone_hsdp_basic(self):
|
||||
"""Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 1
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 4
|
||||
assert result["effective_dp_size"] == 4
|
||||
|
||||
def test_standalone_hsdp_with_replicate(self):
|
||||
"""Test standalone HSDP with replication (shard_size=4, replicate=2)."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=2,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 1
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 8 # 4 * 2
|
||||
assert result["effective_dp_size"] == 8
|
||||
|
||||
def test_combined_hsdp_sp_not_standalone(self):
|
||||
"""Test HSDP combined with SP is NOT detected as standalone.
|
||||
|
||||
This is a regression test for the bug where the condition
|
||||
`dit_parallel_size == fully_shard_degree` incorrectly matched
|
||||
combined modes like SP=4 + HSDP=4.
|
||||
"""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
sequence_parallel_size=4,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 4
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
# Should NOT override dp_size for combined mode
|
||||
assert result["effective_dp_size"] == 1 # original data_parallel_size
|
||||
|
||||
def test_combined_hsdp_cfg_not_standalone(self):
|
||||
"""Test HSDP combined with CFG is NOT detected as standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
cfg_parallel_size=2,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 2
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_combined_hsdp_pp_not_standalone(self):
|
||||
"""Test HSDP combined with PP is NOT detected as standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
pipeline_parallel_size=2,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 2
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_no_hsdp_not_standalone(self):
|
||||
"""Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 1
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_combined_multiple_parallelism_not_standalone(self):
|
||||
"""Test HSDP combined with multiple parallelism is NOT standalone."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
sequence_parallel_size=2,
|
||||
cfg_parallel_size=2,
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["original_dit_parallel_size"] == 4 # 2 * 2
|
||||
assert result["is_standalone_hsdp"] is False
|
||||
assert result["effective_dp_size"] == 1
|
||||
|
||||
def test_standalone_hsdp_large_shard(self):
|
||||
"""Test standalone HSDP with large shard size."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=8,
|
||||
hsdp_replicate_size=1,
|
||||
)
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 8
|
||||
assert result["effective_dp_size"] == 8
|
||||
|
||||
def test_standalone_hsdp_large_replicate(self):
|
||||
"""Test standalone HSDP with large replicate size."""
|
||||
result = self.compute_standalone_hsdp_params(
|
||||
fully_shard_degree=4,
|
||||
hsdp_replicate_size=4,
|
||||
)
|
||||
assert result["is_standalone_hsdp"] is True
|
||||
assert result["effective_dit_parallel_size"] == 16 # 4 * 4
|
||||
assert result["effective_dp_size"] == 16
|
||||
|
||||
|
||||
class TestHSDPShardConditions:
|
||||
"""Tests for _hsdp_shard_conditions matching logic."""
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ from huggingface_hub import snapshot_download
|
||||
from vllm.config.load import LoadConfig
|
||||
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
||||
HostWeightPlan,
|
||||
@@ -91,17 +91,6 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
|
||||
return loader
|
||||
|
||||
|
||||
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model="",
|
||||
dtype=torch.float32,
|
||||
quantization_config="fp8",
|
||||
parallel_config=DiffusionParallelConfig(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
|
||||
|
||||
def test_strict_check_only_validates_source_prefix_parameters():
|
||||
model = _DummyPipelineModel(source_prefix="transformer.")
|
||||
loader = _make_loader_with_weights(["transformer.weight"])
|
||||
@@ -405,78 +394,6 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerTensorOnlineLinearMethod,
|
||||
)
|
||||
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
|
||||
od_config = _make_dlo_online_quant_config()
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
|
||||
model.transformer.quant_method.uses_meta_device = True
|
||||
calls: list[object] = []
|
||||
allowlist_models: list[nn.Module] = []
|
||||
|
||||
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
|
||||
|
||||
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
|
||||
allowlist_models.append(candidate)
|
||||
return original_allowlist_check(candidate)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
|
||||
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
|
||||
loader.load_weights = ( # type: ignore[method-assign]
|
||||
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
|
||||
)
|
||||
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(
|
||||
None,
|
||||
"online quantization requires the ordinary loader",
|
||||
),
|
||||
)
|
||||
|
||||
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
|
||||
assert allowlist_models == [model]
|
||||
assert calls == [("load", True), "process"]
|
||||
assert loader.take_host_weight_plan() is None
|
||||
|
||||
|
||||
def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
|
||||
class UnsupportedOnlineMethod:
|
||||
uses_meta_device = True
|
||||
|
||||
od_config = _make_dlo_online_quant_config()
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
model.transformer.quant_method = UnsupportedOnlineMethod()
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
monkeypatch.setattr(
|
||||
loader_mod,
|
||||
"build_checkpoint_mmap_plan",
|
||||
lambda *_args, **_kwargs: HostWeightPlanResult(
|
||||
None,
|
||||
"online quantization requires the ordinary loader",
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="per-tensor FP8 linears"):
|
||||
loader.load_model(load_device="cpu")
|
||||
|
||||
|
||||
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
|
||||
|
||||
@@ -118,14 +118,6 @@ def _call_mgr(
|
||||
)
|
||||
|
||||
|
||||
def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
|
||||
mgr = _make_cache_mgr()
|
||||
|
||||
assert isinstance(mgr, nn.Module)
|
||||
assert dict(mgr.named_modules())["attn"] is mgr.attn
|
||||
assert mgr.state_dict() == {}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test 1: No AR KV — basic cache → reuse
|
||||
# ============================================================
|
||||
|
||||
Whitespace-only changes.
@@ -1,129 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Equivalence tests for SANA-WM's camera-branch reduce-scatter.
|
||||
|
||||
``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the
|
||||
full hidden width, while the GDN stream it is added to is TP-local. Instead of
|
||||
all-reducing and then discarding all but one slice, the branch reduce-scatters
|
||||
straight to the local shard. These tests pin the shard each rank ends up with,
|
||||
including the bias slice, on both code paths:
|
||||
|
||||
* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor``
|
||||
directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement
|
||||
``reduce_scatter``;
|
||||
* the all-reduce-and-slice fallback taken when the coordinator exposes no
|
||||
``device_group``.
|
||||
|
||||
CPU-only: the collectives and the rank/world-size lookups are mocked, so a
|
||||
single process can simulate an arbitrary TP degree.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
|
||||
|
||||
HIDDEN = 24
|
||||
|
||||
|
||||
def _host(bias: torch.Tensor | None) -> SimpleNamespace:
|
||||
"""A stand-in ``self`` exposing only what the method reads."""
|
||||
return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias))
|
||||
|
||||
|
||||
def _fake_reduce_scatter_tensor(rank: int, tp_size: int):
|
||||
"""Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``.
|
||||
|
||||
The test drives a single process, so the other ranks' partials are already
|
||||
folded into the input; scattering is all this has to reproduce.
|
||||
"""
|
||||
|
||||
def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None:
|
||||
output.copy_(input_.chunk(tp_size, dim=0)[rank])
|
||||
|
||||
return _impl
|
||||
|
||||
|
||||
def _tp_group(with_device_group: bool):
|
||||
return SimpleNamespace(device_group=object() if with_device_group else None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp_size", [2, 4])
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None:
|
||||
"""The scattered shard equals the slice the all-reduce path would take."""
|
||||
local_width = HIDDEN // tp_size
|
||||
torch.manual_seed(0)
|
||||
contrib = torch.randn(2, 3, HIDDEN)
|
||||
bias = torch.randn(HIDDEN)
|
||||
reference = torch.zeros(2, 3, local_width)
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
|
||||
patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
|
||||
patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)),
|
||||
patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)),
|
||||
):
|
||||
out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
|
||||
|
||||
start = rank * local_width
|
||||
expected = (contrib + bias)[..., start : start + local_width]
|
||||
assert out.shape == (2, 3, local_width)
|
||||
torch.testing.assert_close(out, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
def test_fallback_path_matches_the_process_group_path(rank: int) -> None:
|
||||
"""Without a device_group the method all-reduces and slices, same result."""
|
||||
tp_size = 2
|
||||
local_width = HIDDEN // tp_size
|
||||
torch.manual_seed(1)
|
||||
contrib = torch.randn(2, 3, HIDDEN)
|
||||
bias = torch.randn(HIDDEN)
|
||||
reference = torch.zeros(2, 3, local_width)
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
|
||||
patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
|
||||
patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)),
|
||||
patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t),
|
||||
):
|
||||
out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
|
||||
|
||||
start = rank * local_width
|
||||
torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width])
|
||||
|
||||
|
||||
def test_tp1_adds_bias_without_a_collective() -> None:
|
||||
contrib = torch.randn(2, 3, HIDDEN)
|
||||
bias = torch.randn(HIDDEN)
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1),
|
||||
patch("torch.distributed.reduce_scatter_tensor") as collective,
|
||||
):
|
||||
out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib)
|
||||
|
||||
assert collective.call_count == 0
|
||||
torch.testing.assert_close(out, contrib + bias)
|
||||
|
||||
|
||||
def test_rejects_a_width_that_does_not_shard() -> None:
|
||||
"""A local width that is not 1/tp of the contribution is a wiring bug."""
|
||||
contrib = torch.randn(1, 2, HIDDEN)
|
||||
reference = torch.zeros(1, 2, HIDDEN // 3) # tp=2 would need HIDDEN // 2
|
||||
|
||||
with (
|
||||
patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
|
||||
patch("torch.distributed.reduce_scatter_tensor"),
|
||||
pytest.raises(ValueError, match="camera contribution width mismatch"),
|
||||
):
|
||||
SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference)
|
||||
@@ -1,111 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion.
|
||||
|
||||
`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
|
||||
and issues a single tensor-parallel all-reduce instead of two. These tests lock
|
||||
in that it is numerically identical to applying `norm_q`/`norm_k` separately,
|
||||
and that it really issues exactly one (fused) all-reduce.
|
||||
|
||||
CPU-only: the collective is mocked, so a single process simulates one rank
|
||||
(all-reduce is identity). This exercises the cat/split logic a real multi-rank
|
||||
run depends on.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
|
||||
TensorParallelRMSNorm,
|
||||
fused_qk_rms_norm,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
|
||||
|
||||
|
||||
def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]:
|
||||
"""Two norms with distinct random weights so a swapped q/k slice is caught."""
|
||||
torch.manual_seed(0)
|
||||
norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size)
|
||||
norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size)
|
||||
norm_q.weight.data.normal_()
|
||||
norm_k.weight.data.normal_()
|
||||
return norm_q, norm_k
|
||||
|
||||
|
||||
def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Single-rank all-reduce: the reduction over one rank is the identity."""
|
||||
return tensor
|
||||
|
||||
|
||||
def test_fused_matches_separate() -> None:
|
||||
dim = 64
|
||||
norm_q, norm_k = _make_norms(dim, tp_size=2)
|
||||
q = torch.randn(2, 5, dim)
|
||||
k = torch.randn(2, 5, dim)
|
||||
|
||||
with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce):
|
||||
ref_q, ref_k = norm_q(q), norm_k(k)
|
||||
fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
|
||||
|
||||
torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
|
||||
torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_fused_issues_one_all_reduce() -> None:
|
||||
dim = 32
|
||||
norm_q, norm_k = _make_norms(dim, tp_size=4)
|
||||
q = torch.randn(1, 3, dim)
|
||||
k = torch.randn(1, 3, dim)
|
||||
|
||||
with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
|
||||
fused_qk_rms_norm(norm_q, norm_k, q, k)
|
||||
assert collective.call_count == 1
|
||||
# Both sums travel in one tensor: last dim is 2, not 1.
|
||||
assert collective.call_args.args[0].shape[-1] == 2
|
||||
|
||||
collective.reset_mock()
|
||||
norm_q(q), norm_k(k)
|
||||
assert collective.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"norms",
|
||||
[
|
||||
(RMSNorm(16), RMSNorm(16)),
|
||||
(nn.Identity(), nn.Identity()),
|
||||
],
|
||||
ids=["plain_rms_norm", "identity"],
|
||||
)
|
||||
def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None:
|
||||
"""qk_norm=False and TP=1 both hand in norms the fused path cannot use."""
|
||||
norm_q, norm_k = norms
|
||||
q = torch.randn(2, 4, 16)
|
||||
k = torch.randn(2, 4, 16)
|
||||
|
||||
with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
|
||||
fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
|
||||
assert collective.call_count == 0
|
||||
|
||||
torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
|
||||
torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_tp1_skips_the_collective() -> None:
|
||||
norm_q, norm_k = _make_norms(16, tp_size=1)
|
||||
q = torch.randn(2, 4, 16)
|
||||
k = torch.randn(2, 4, 16)
|
||||
|
||||
with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
|
||||
fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
|
||||
assert collective.call_count == 0
|
||||
|
||||
torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
|
||||
torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
|
||||
@@ -1,142 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for SANA-WM request normalization."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.sana_wm.request import (
|
||||
SANA_WM_DEFAULT_HEIGHT,
|
||||
SANA_WM_DEFAULT_NUM_FRAMES,
|
||||
SANA_WM_DEFAULT_WIDTH,
|
||||
normalize_sana_wm_payload,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _prompt(**sana_wm):
|
||||
"""A valid request, with the action rollout kept in step with num_frames."""
|
||||
block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm}
|
||||
block.setdefault("action", f"w-{block['num_frames'] - 1}")
|
||||
return {
|
||||
"prompt": "a quiet city street",
|
||||
"multi_modal_data": {"image": object()},
|
||||
"sana_wm": block,
|
||||
}
|
||||
|
||||
|
||||
def _identity_poses(count):
|
||||
return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)]
|
||||
|
||||
|
||||
def _camera_prompt(**camera):
|
||||
prompt = _prompt(num_frames=9)
|
||||
prompt["sana_wm"].pop("action")
|
||||
prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera}
|
||||
return prompt
|
||||
|
||||
|
||||
def test_defaults_are_latent_aligned():
|
||||
payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"]
|
||||
assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES
|
||||
assert payload["height"] == SANA_WM_DEFAULT_HEIGHT
|
||||
assert payload["width"] == SANA_WM_DEFAULT_WIDTH
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_frames", [9, 33, 161])
|
||||
def test_accepts_aligned_geometry(num_frames):
|
||||
payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"]
|
||||
assert payload["num_frames"] == num_frames
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_frames", [8, 10, 160])
|
||||
def test_rejects_unaligned_num_frames(num_frames):
|
||||
# The VAE compresses 8x temporally; an unaligned count would be floored.
|
||||
with pytest.raises(ValueError, match="num_frames"):
|
||||
normalize_sana_wm_payload(_prompt(num_frames=num_frames))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)])
|
||||
def test_rejects_unaligned_resolution(height, width):
|
||||
with pytest.raises(ValueError, match="divisible by 32"):
|
||||
normalize_sana_wm_payload(_prompt(height=height, width=width))
|
||||
|
||||
|
||||
def test_normalization_is_idempotent():
|
||||
once = normalize_sana_wm_payload(_prompt(num_frames=33))
|
||||
twice = normalize_sana_wm_payload(once)
|
||||
assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"]
|
||||
|
||||
|
||||
def test_rejects_action_shorter_than_num_frames():
|
||||
# Used to shrink the camera tensor silently, then fail deep in the
|
||||
# transformer on a camera/latent frame mismatch.
|
||||
with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"):
|
||||
normalize_sana_wm_payload(_prompt(action="w-32"))
|
||||
|
||||
|
||||
def test_rejects_action_longer_than_num_frames():
|
||||
with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"):
|
||||
normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160"))
|
||||
|
||||
|
||||
def test_rejects_oversized_action_without_expanding_it():
|
||||
# A 12-byte request body: expanding this rollout before checking the frame
|
||||
# count allocated ~16 GB. Rejection has to be arithmetic, so this test also
|
||||
# fails (by hanging/OOM) if the duration is ever expanded again.
|
||||
with pytest.raises(ValueError, match="rolls out 1000000001 frames"):
|
||||
normalize_sana_wm_payload(_prompt(action="w-1000000000"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("action", ["", " ", "w-0", "w-abc", "w", "x-4"])
|
||||
def test_rejects_malformed_action(action):
|
||||
with pytest.raises(ValueError):
|
||||
normalize_sana_wm_payload(_prompt(action=action))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")])
|
||||
def test_rejects_unusable_focal_length(value):
|
||||
# compute_raymap divides by fx/fy; a zero or non-finite focal length turns
|
||||
# the whole Plucker map into NaN and the response into a NaN video.
|
||||
intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352}
|
||||
with pytest.raises(ValueError, match="fx"):
|
||||
normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf")])
|
||||
def test_rejects_non_finite_principal_point(value):
|
||||
intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352}
|
||||
with pytest.raises(ValueError, match="must be finite"):
|
||||
normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf")])
|
||||
def test_rejects_non_finite_speeds(value):
|
||||
# `parsed <= 0` is False for NaN, so the positivity check alone let it pass.
|
||||
with pytest.raises(ValueError, match="must be finite"):
|
||||
normalize_sana_wm_payload(_prompt(translation_speed=value))
|
||||
|
||||
|
||||
def test_rejects_non_finite_pose_entry():
|
||||
prompt = _camera_prompt()
|
||||
prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan")
|
||||
with pytest.raises(ValueError, match="must be finite"):
|
||||
normalize_sana_wm_payload(prompt)
|
||||
|
||||
|
||||
def test_accepts_default_camera_metadata():
|
||||
payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"]
|
||||
assert payload["camera"]["poses"]
|
||||
# The unread fields are not carried into the canonical payload.
|
||||
assert set(payload["camera"]) == {"poses"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[("format", "w2c_4x4"), ("coordinate_system", "opengl")],
|
||||
)
|
||||
def test_rejects_unsupported_camera_convention(field, value):
|
||||
# Nothing downstream reads these, so accepting one would mean interpreting
|
||||
# the poses under the wrong convention.
|
||||
with pytest.raises(ValueError, match=field.split("_")[0]):
|
||||
normalize_sana_wm_payload(_camera_prompt(**{field: value}))
|
||||
@@ -1,72 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion
|
||||
config shim must not crash on missing ``head_dtype`` and the resulting
|
||||
LogitsProcessor must have ``head_dtype is None`` (= use model dtype)."""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.parallel_state import (
|
||||
cleanup_dist_env_and_memory,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
|
||||
SenseNovaU1ForCausalLM,
|
||||
)
|
||||
from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
|
||||
from vllm_omni.transformers_utils.configs.sensenova_u1 import (
|
||||
SenseNovaU1Config,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _build_tiny_sensenova_u1(quant_config=None, prefix="model"):
|
||||
"""Build a minimal SenseNovaU1ForCausalLM for unit tests."""
|
||||
config = SenseNovaU1Config(
|
||||
llm_config={
|
||||
"hidden_size": 64,
|
||||
"num_attention_heads": 2,
|
||||
"num_hidden_layers": 1,
|
||||
"intermediate_size": 128,
|
||||
"vocab_size": 32,
|
||||
"max_position_embeddings": 128,
|
||||
"max_position_embeddings_hw": 128,
|
||||
},
|
||||
)
|
||||
return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix)
|
||||
|
||||
|
||||
def test_logits_processor_head_dtype_under_diffusion_shim():
|
||||
# Initialize the distributed environment
|
||||
os.environ.setdefault("MASTER_ADDR", "localhost")
|
||||
os.environ.setdefault("MASTER_PORT", "29543")
|
||||
init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
|
||||
initialize_model_parallel()
|
||||
|
||||
# Create a vLLM config with a wrapped diffusion config
|
||||
fake_diff_config = _DiffusionVllmModelConfig(
|
||||
model="sensenova-u1-test",
|
||||
dtype=torch.bfloat16,
|
||||
max_model_len=8192,
|
||||
original_max_model_len=8192,
|
||||
)
|
||||
vllm_config = VllmConfig(
|
||||
device_config=DeviceConfig(device="cpu"),
|
||||
)
|
||||
vllm_config.modelconfig = fake_diff_config # type: ignore[assignment]
|
||||
|
||||
# Build a tiny version of the Causal LM component
|
||||
with set_current_vllm_config(vllm_config):
|
||||
model = _build_tiny_sensenova_u1()
|
||||
|
||||
# Ensure that we have a set head_dtype attribute. Currently, the head_dtype
|
||||
# is set to None
|
||||
head_dtype = model.logits_processor.head_dtype
|
||||
assert head_dtype is None
|
||||
cleanup_dist_env_and_memory()
|
||||
@@ -185,79 +185,6 @@ class TestDistributedLayerwiseOffloadHook:
|
||||
assert next_block.weight.stride() == expected_stride
|
||||
assert torch.equal(next_block.weight, expected)
|
||||
|
||||
def test_allgather_reconstructs_online_fp8_weight_and_scale(
|
||||
self,
|
||||
patched_offload_runtime,
|
||||
monkeypatch,
|
||||
):
|
||||
"""DLO must reconstruct finalized online-FP8 tensors byte-for-byte."""
|
||||
|
||||
class OnlineFp8Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
logical_weight = torch.tensor(
|
||||
[
|
||||
[1.0, 2.0, 3.0, 4.0],
|
||||
[5.0, 6.0, 7.0, 8.0],
|
||||
[9.0, 10.0, 11.0, 12.0],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
).to(torch.float8_e4m3fn)
|
||||
self.weight = nn.Parameter(logical_weight.t(), requires_grad=False)
|
||||
self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False)
|
||||
|
||||
current_block = OnlineFp8Block()
|
||||
rank0_block = OnlineFp8Block()
|
||||
rank1_block = OnlineFp8Block()
|
||||
expected_weight = rank0_block.weight.detach().clone()
|
||||
expected_scale = rank0_block.weight_scale.detach().clone()
|
||||
expected_stride = rank0_block.weight.stride()
|
||||
|
||||
rank0_hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=rank0_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=object(),
|
||||
dp_size=2,
|
||||
rank=0,
|
||||
copy_stream=DummyStream(),
|
||||
comm_stream=DummyStream(),
|
||||
pin_memory=False,
|
||||
)
|
||||
rank1_hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=rank1_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=object(),
|
||||
dp_size=2,
|
||||
rank=1,
|
||||
copy_stream=DummyStream(),
|
||||
comm_stream=DummyStream(),
|
||||
pin_memory=False,
|
||||
)
|
||||
rank0_hook.initialize_hook(current_block)
|
||||
rank1_hook.initialize_hook(OnlineFp8Block())
|
||||
rank0_hook.gpu_shard_buffers = [
|
||||
{dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
|
||||
]
|
||||
|
||||
def fake_allgather(output, local_shard, *, group):
|
||||
del group
|
||||
remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
|
||||
shard_size = local_shard.numel()
|
||||
output[:shard_size].copy_(local_shard)
|
||||
output[shard_size : 2 * shard_size].copy_(remote_shard)
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
|
||||
rank0_hook.prefetch_layer(slot=0, non_blocking=False)
|
||||
|
||||
assert rank0_block.weight.stride() == expected_stride
|
||||
torch.testing.assert_close(
|
||||
rank0_block.weight.float(),
|
||||
expected_weight.float(),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
|
||||
|
||||
def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
|
||||
"""DTensor wrapper should be preserved through prefetch/offload cycle."""
|
||||
current_block = TinyBlock(_make_values(1.0))
|
||||
|
||||
@@ -92,30 +92,6 @@ class TestParallelConfigPropagation:
|
||||
od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
|
||||
assert od.num_gpus == 2
|
||||
|
||||
def test_dp_is_inferred_from_num_gpus(self):
|
||||
pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2)
|
||||
od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8)
|
||||
assert od.parallel_config.data_parallel_size == 2
|
||||
assert od.parallel_config.world_size == 8
|
||||
|
||||
def test_explicit_dp_is_validated_against_num_gpus(self):
|
||||
pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2)
|
||||
od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
|
||||
assert od.parallel_config.data_parallel_size == 2
|
||||
|
||||
with pytest.raises(ValueError, match="does not match WORLD-derived value"):
|
||||
OmniDiffusionConfig.from_kwargs(
|
||||
model="x",
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2),
|
||||
num_gpus=8,
|
||||
)
|
||||
|
||||
def test_hsdp_does_not_infer_ordinary_dp(self):
|
||||
pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4)
|
||||
od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
|
||||
assert od.parallel_config.data_parallel_size == 1
|
||||
assert od.parallel_config.world_size == 4
|
||||
|
||||
|
||||
class TestCreateDefaultDiffusion:
|
||||
"""Verify engine_args structure from create_default_diffusion."""
|
||||
@@ -154,21 +130,6 @@ class TestCreateDefaultDiffusion:
|
||||
|
||||
assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
|
||||
def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
|
||||
od = _roundtrip_diffusion_config(
|
||||
model="x",
|
||||
kv_cache_memory_bytes=4096,
|
||||
gpu_memory_utilization=0.75,
|
||||
max_num_batched_tokens=2048,
|
||||
max_model_len=4096,
|
||||
)
|
||||
|
||||
assert od.kv_cache_memory_bytes == 4096
|
||||
assert od.gpu_memory_utilization == 0.75
|
||||
assert od.max_num_batched_tokens == 2048
|
||||
assert od.max_model_len == 4096
|
||||
|
||||
|
||||
def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
|
||||
od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
|
||||
|
||||
@@ -161,83 +161,9 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa
|
||||
assert exc_info.value is initialization_error
|
||||
custom_scheduler.initialize.assert_called_once_with(od_config)
|
||||
custom_scheduler.close.assert_called_once_with()
|
||||
|
||||
|
||||
def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
fake_executor = SimpleNamespace(shutdown=Mock())
|
||||
monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
|
||||
monkeypatch.setattr(
|
||||
DiffusionEngine,
|
||||
"_resolve_execution_mode",
|
||||
lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
|
||||
Mock(side_effect=RuntimeError("KV initialization failed")),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="KV initialization failed"):
|
||||
DiffusionEngine(SimpleNamespace())
|
||||
|
||||
fake_executor.shutdown.assert_called_once_with()
|
||||
|
||||
|
||||
def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
events = []
|
||||
expected_profile_requests = [object()]
|
||||
fake_executor = SimpleNamespace(shutdown=Mock())
|
||||
kv_cache_config = object()
|
||||
kv_vllm_config = object()
|
||||
od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")
|
||||
|
||||
monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
|
||||
monkeypatch.setattr(
|
||||
DiffusionEngine,
|
||||
"_resolve_execution_mode",
|
||||
lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DiffusionEngine,
|
||||
"_init_executor",
|
||||
lambda self, config: setattr(self, "executor", fake_executor),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
DiffusionEngine,
|
||||
"_prepare_diffusion_kv_profile_requests",
|
||||
lambda self: events.append("prepare-profile") or expected_profile_requests,
|
||||
)
|
||||
|
||||
def initialize(executor, config, *, profile_requests):
|
||||
events.append("profile-workers")
|
||||
assert executor is fake_executor
|
||||
assert config is od_config
|
||||
assert profile_requests is expected_profile_requests
|
||||
return kv_cache_config, 16, 16, kv_vllm_config
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
|
||||
initialize,
|
||||
)
|
||||
|
||||
def init_scheduler(self, config, scheduler, *args, **kwargs):
|
||||
events.append("initialize-scheduler")
|
||||
self.scheduler = SimpleNamespace(close=Mock())
|
||||
|
||||
monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
|
||||
monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
|
||||
monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
|
||||
monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)
|
||||
|
||||
DiffusionEngine(od_config)
|
||||
|
||||
assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_step_compatibility_wrapper_returns_final_batch() -> None:
|
||||
engine = _make_engine()
|
||||
|
||||
@@ -1,17 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion import io_support
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
@@ -63,151 +55,3 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk
|
||||
)
|
||||
|
||||
assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
|
||||
|
||||
|
||||
def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
|
||||
assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
|
||||
assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
|
||||
|
||||
|
||||
def test_dense_mode_does_not_build_kv_profile_request() -> None:
|
||||
engine = object.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
|
||||
engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))
|
||||
|
||||
assert engine._prepare_diffusion_kv_profile_requests() is None
|
||||
engine._make_dummy_request.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
|
||||
[
|
||||
(DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
|
||||
(DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
|
||||
(DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
|
||||
],
|
||||
)
|
||||
def test_paged_kv_profile_requests_match_per_rank_batch(
|
||||
execution_mode: DiffusionExecutionMode,
|
||||
uses_dlo_dp: bool,
|
||||
max_num_seqs: int,
|
||||
expected_profile_requests: int,
|
||||
) -> None:
|
||||
engine = object.__new__(DiffusionEngine)
|
||||
engine.execution_mode = execution_mode
|
||||
engine.od_config = SimpleNamespace(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
model_class_name="HunyuanImage3ForCausalMM",
|
||||
max_num_seqs=max_num_seqs,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
|
||||
enable_distributed_layerwise_offload=uses_dlo_dp,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
request = OmniDiffusionRequest(
|
||||
prompt="profile",
|
||||
request_id="profile-request",
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
)
|
||||
engine._make_dummy_request = Mock(return_value=request)
|
||||
prepared_layout = object()
|
||||
scheduler_kv_state = (
|
||||
DiffusionKVRequest(
|
||||
"profile-request/diffusion-kv/0",
|
||||
sequence_id=0,
|
||||
prefix_len=3,
|
||||
target_len=8,
|
||||
seq_len=12,
|
||||
),
|
||||
DiffusionKVRequest(
|
||||
"profile-request/diffusion-kv/1",
|
||||
sequence_id=1,
|
||||
prefix_len=4,
|
||||
target_len=8,
|
||||
seq_len=13,
|
||||
),
|
||||
)
|
||||
|
||||
def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
|
||||
req.prepared_layout = prepared_layout
|
||||
req.diffusion_kv_requests = scheduler_kv_state
|
||||
return req
|
||||
|
||||
engine._prepare_request_for_admission = Mock(side_effect=preprocess)
|
||||
|
||||
result = engine._prepare_diffusion_kv_profile_requests()
|
||||
|
||||
assert result is not None
|
||||
assert len(result) == expected_profile_requests
|
||||
assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
|
||||
assert all(profile_request.is_dummy_run() for profile_request in result)
|
||||
assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
|
||||
assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
|
||||
assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
|
||||
assert engine._diffusion_kv_profile_limits == (2, 13, 8)
|
||||
assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
|
||||
engine._make_dummy_request.assert_called_once_with(
|
||||
height=1024,
|
||||
width=1024,
|
||||
guidance_scale=5.0,
|
||||
num_image_inputs=3,
|
||||
)
|
||||
engine._prepare_request_for_admission.assert_called_once_with(request)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("num_sequences", "seq_len", "target_len"),
|
||||
[
|
||||
(3, 16, 8),
|
||||
(1, 17, 8),
|
||||
(1, 16, 9),
|
||||
],
|
||||
)
|
||||
def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
|
||||
num_sequences: int,
|
||||
seq_len: int,
|
||||
target_len: int,
|
||||
) -> None:
|
||||
engine = object.__new__(DiffusionEngine)
|
||||
engine._diffusion_kv_profile_limits = (2, 16, 8)
|
||||
engine.pre_process_func = lambda request: request
|
||||
request = OmniDiffusionRequest(
|
||||
prompt="too large",
|
||||
request_id="too-large",
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
diffusion_kv_requests=tuple(
|
||||
DiffusionKVRequest(
|
||||
f"too-large/diffusion-kv/{sequence_id}",
|
||||
sequence_id=sequence_id,
|
||||
prefix_len=0,
|
||||
target_len=target_len,
|
||||
seq_len=seq_len,
|
||||
)
|
||||
for sequence_id in range(num_sequences)
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
|
||||
engine._prepare_request_for_admission(request)
|
||||
|
||||
|
||||
def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
|
||||
engine = object.__new__(DiffusionEngine)
|
||||
engine._diffusion_kv_profile_limits = (2, 16, 8)
|
||||
engine.pre_process_func = lambda request: request
|
||||
request = OmniDiffusionRequest(
|
||||
prompt="fits",
|
||||
request_id="fits",
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
diffusion_kv_requests=tuple(
|
||||
DiffusionKVRequest(
|
||||
f"fits/diffusion-kv/{sequence_id}",
|
||||
sequence_id=sequence_id,
|
||||
prefix_len=8,
|
||||
target_len=8,
|
||||
seq_len=16,
|
||||
)
|
||||
for sequence_id in range(2)
|
||||
),
|
||||
)
|
||||
|
||||
assert engine._prepare_request_for_admission(request) is request
|
||||
@@ -473,38 +473,6 @@ async def test_cancelled_future_does_not_kill_busy_loop():
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
|
||||
loop = asyncio.get_running_loop()
|
||||
engine = _make_engine_with_loop(loop)
|
||||
|
||||
def fail_reservation(*args, **kwargs):
|
||||
raise RuntimeError("native allocation bug")
|
||||
|
||||
engine.scheduler._diffusion_kv_manager = SimpleNamespace(
|
||||
has_request=lambda request_id: False,
|
||||
reserve_request=fail_reservation,
|
||||
free_request=lambda request_id: None,
|
||||
)
|
||||
request = _make_request("kv-error")
|
||||
request.diffusion_kv_requests = (object(),)
|
||||
try:
|
||||
response_stream = engine.async_add_req_and_stream_response(request)
|
||||
output_queue = engine._out_streams[request.request_id]
|
||||
for _ in range(300):
|
||||
if not output_queue.empty():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"
|
||||
|
||||
output = await anext(response_stream)
|
||||
assert output.error == "native allocation bug"
|
||||
assert output.finished
|
||||
assert engine.worker_thread.is_alive()
|
||||
finally:
|
||||
_stop_engine(engine)
|
||||
|
||||
|
||||
# ─────────────────────────── _RpcTask basics ───────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -513,84 +512,6 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
|
||||
assert runner.pipeline.last_req.num_reqs == 1
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch):
|
||||
runner = _make_runner(cache_backend=None, cache_backend_name="none")
|
||||
request = _make_request()
|
||||
runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
|
||||
record_names = []
|
||||
original_execute = runner._execute_request_list
|
||||
|
||||
def execute(*args, **kwargs):
|
||||
record_names.append(kwargs["record_name"])
|
||||
return original_execute(*args, **kwargs)
|
||||
|
||||
runner._execute_request_list = execute
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
reset_peak_memory_stats = Mock()
|
||||
monkeypatch.setattr(
|
||||
model_runner_module.current_omni_platform,
|
||||
"reset_peak_memory_stats",
|
||||
reset_peak_memory_stats,
|
||||
)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
||||
synchronize = Mock()
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
|
||||
|
||||
runner.profile_run([request])
|
||||
|
||||
assert runner.pipeline.forward_calls == 1
|
||||
assert record_names == ["pipeline_memory_profile"]
|
||||
runner._validate_diffusion_kv_metadata.assert_not_called()
|
||||
reset_peak_memory_stats.assert_not_called()
|
||||
synchronize.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch):
|
||||
runner = _make_runner(cache_backend=None, cache_backend_name=None)
|
||||
runner.pipeline = _FinalOnlyStepPipeline()
|
||||
runner.od_config.step_execution = True
|
||||
requests = [_make_request(), _make_request()]
|
||||
requests[0].request_id = "profile-0"
|
||||
requests[1].request_id = "profile-1"
|
||||
for request in requests:
|
||||
request.sampling_params.num_inference_steps = 1
|
||||
|
||||
observed_batch_rows = []
|
||||
original_denoise_step = runner.pipeline.denoise_step
|
||||
|
||||
def denoise_step(input_batch, states):
|
||||
observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
|
||||
return original_denoise_step(input_batch, states)
|
||||
|
||||
runner.pipeline.denoise_step = denoise_step
|
||||
runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
reset_peak_memory_stats = Mock()
|
||||
monkeypatch.setattr(
|
||||
model_runner_module.current_omni_platform,
|
||||
"reset_peak_memory_stats",
|
||||
reset_peak_memory_stats,
|
||||
)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
|
||||
synchronize = Mock()
|
||||
monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
|
||||
|
||||
runner.profile_run(requests)
|
||||
|
||||
assert observed_batch_rows == [(2, 2)]
|
||||
assert runner.state_cache == {}
|
||||
assert runner.input_batch is None
|
||||
runner._validate_diffusion_kv_metadata.assert_not_called()
|
||||
reset_peak_memory_stats.assert_not_called()
|
||||
synchronize.assert_called_once_with()
|
||||
|
||||
|
||||
class _BatchPipeline:
|
||||
"""Pipeline returning a configurable list of outputs from forward()."""
|
||||
|
||||
|
||||
@@ -52,7 +52,6 @@ def test_dummy_run_request_is_identified_by_reserved_request_id():
|
||||
|
||||
assert req.is_dummy_run()
|
||||
assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
|
||||
assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1")
|
||||
|
||||
|
||||
def test_non_dummy_request_is_not_identified_as_dummy_run():
|
||||
|
||||
@@ -8,14 +8,12 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
|
||||
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched import (
|
||||
BaseScheduler,
|
||||
@@ -89,52 +87,6 @@ def _cached_ids(sched_output) -> list[str]:
|
||||
return list(sched_output.scheduled_cached_reqs.request_ids)
|
||||
|
||||
|
||||
def _initialize_paged_scheduler(
|
||||
scheduler: BaseScheduler,
|
||||
*,
|
||||
num_blocks: int = 64,
|
||||
max_num_seqs: int = 1,
|
||||
) -> None:
|
||||
native_kv_managers.register_all_kvcache_specs(None)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=4,
|
||||
num_kv_heads=2,
|
||||
head_size=8,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
|
||||
kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
|
||||
)
|
||||
scheduler.initialize(
|
||||
SimpleNamespace(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
max_model_len=64,
|
||||
max_num_seqs=max_num_seqs,
|
||||
),
|
||||
kv_cache_config=config,
|
||||
scheduler_block_size=4,
|
||||
hash_block_size=4,
|
||||
kv_vllm_config=SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=64),
|
||||
max_in_flight_tokens=64,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
|
||||
request.diffusion_kv_requests = (
|
||||
DiffusionKVRequest(
|
||||
f"{request.request_id}/diffusion-kv/0",
|
||||
sequence_id=0,
|
||||
prefix_len=4,
|
||||
target_len=4,
|
||||
seq_len=seq_len,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _StubScheduler:
|
||||
def __init__(self, request: OmniDiffusionRequest, output) -> None:
|
||||
self._request = request
|
||||
@@ -448,16 +400,18 @@ class TestRequestScheduler:
|
||||
assert self.scheduler.has_requests() is False
|
||||
|
||||
def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler)
|
||||
self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
|
||||
request = _make_request("diffusion-kv")
|
||||
prepared_layout = object()
|
||||
request.prepared_layout = prepared_layout
|
||||
text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
|
||||
kv_request = DiffusionKVRequest(
|
||||
"diffusion-kv/diffusion-kv/0",
|
||||
sequence_id=0,
|
||||
prefix_len=4,
|
||||
target_len=8,
|
||||
seq_len=16,
|
||||
kv_contexts=(text_context,),
|
||||
)
|
||||
request.diffusion_kv_requests = (kv_request,)
|
||||
|
||||
@@ -467,179 +421,12 @@ class TestRequestScheduler:
|
||||
|
||||
assert state is not None
|
||||
assert state.diffusion_kv_requests == (kv_request,)
|
||||
assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
|
||||
assert request.diffusion_kv_requests is None
|
||||
assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
|
||||
metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata
|
||||
assert metadata is not None
|
||||
assert metadata.request_id == request.request_id
|
||||
assert len(metadata.sequences[0].block_ids[0]) == 4
|
||||
|
||||
def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
|
||||
first_request = _make_request("first")
|
||||
second_request = _make_request("second")
|
||||
_attach_diffusion_kv(first_request)
|
||||
_attach_diffusion_kv(second_request)
|
||||
self.scheduler.add_request(first_request)
|
||||
self.scheduler.add_request(second_request)
|
||||
|
||||
first_output = self.scheduler.schedule()
|
||||
|
||||
assert _new_ids(first_output) == ["first"]
|
||||
assert first_output.num_waiting_reqs == 1
|
||||
|
||||
self.scheduler.update_from_output(first_output, _make_request_output("first"))
|
||||
second_output = self.scheduler.schedule()
|
||||
|
||||
assert _new_ids(second_output) == ["second"]
|
||||
assert second_output.num_waiting_reqs == 0
|
||||
# The same output carries both the release and the newly admitted block
|
||||
# table; the future Worker data plane must process finished ids first.
|
||||
assert second_output.finished_req_ids == {"first"}
|
||||
assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None
|
||||
|
||||
def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2)
|
||||
impossible = _make_request("impossible")
|
||||
schedulable = _make_request("schedulable")
|
||||
_attach_diffusion_kv(impossible, seq_len=8)
|
||||
schedulable.diffusion_kv_requests = (
|
||||
DiffusionKVRequest(
|
||||
"schedulable/diffusion-kv/0",
|
||||
sequence_id=0,
|
||||
prefix_len=0,
|
||||
target_len=4,
|
||||
seq_len=4,
|
||||
),
|
||||
)
|
||||
self.scheduler.add_request(impossible)
|
||||
self.scheduler.add_request(schedulable)
|
||||
|
||||
sched_output = self.scheduler.schedule()
|
||||
|
||||
failed_state = self.scheduler.get_request_state("impossible")
|
||||
assert failed_state is not None
|
||||
assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
|
||||
assert failed_state.error is not None
|
||||
assert "cannot fit even when the block pool is empty" in failed_state.error
|
||||
assert sched_output.finished_req_ids == {"impossible"}
|
||||
assert _new_ids(sched_output) == ["schedulable"]
|
||||
|
||||
finished = self.scheduler.update_from_output(
|
||||
sched_output,
|
||||
_make_request_output("schedulable"),
|
||||
)
|
||||
|
||||
assert finished == {"impossible", "schedulable"}
|
||||
|
||||
def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2)
|
||||
running = _make_request("running")
|
||||
_attach_diffusion_kv(running)
|
||||
self.scheduler.add_request(running)
|
||||
assert _new_ids(self.scheduler.schedule()) == ["running"]
|
||||
|
||||
impossible = _make_request("impossible")
|
||||
impossible.diffusion_kv_requests = tuple(
|
||||
DiffusionKVRequest(
|
||||
f"impossible/diffusion-kv/{sequence_id}",
|
||||
sequence_id=sequence_id,
|
||||
prefix_len=0,
|
||||
target_len=4,
|
||||
seq_len=4,
|
||||
)
|
||||
for sequence_id in range(4)
|
||||
)
|
||||
schedulable = _make_request("schedulable")
|
||||
schedulable.diffusion_kv_requests = (
|
||||
DiffusionKVRequest(
|
||||
"schedulable/diffusion-kv/0",
|
||||
sequence_id=0,
|
||||
prefix_len=0,
|
||||
target_len=4,
|
||||
seq_len=4,
|
||||
),
|
||||
)
|
||||
self.scheduler.add_request(impossible)
|
||||
self.scheduler.add_request(schedulable)
|
||||
|
||||
sched_output = self.scheduler.schedule()
|
||||
|
||||
failed_state = self.scheduler.get_request_state("impossible")
|
||||
assert failed_state is not None
|
||||
assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
|
||||
assert failed_state.error is not None
|
||||
assert "required_blocks=4, available_blocks=3" in failed_state.error
|
||||
assert sched_output.finished_req_ids == {"impossible"}
|
||||
assert _new_ids(sched_output) == ["schedulable"]
|
||||
|
||||
def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler)
|
||||
request = _make_request("native-error")
|
||||
_attach_diffusion_kv(request)
|
||||
self.scheduler.add_request(request)
|
||||
manager = self.scheduler._diffusion_kv_manager
|
||||
assert manager is not None
|
||||
|
||||
def raise_native_error(*args, **kwargs):
|
||||
raise ValueError("native allocation bug")
|
||||
|
||||
monkeypatch.setattr(
|
||||
manager,
|
||||
"reserve_request",
|
||||
raise_native_error,
|
||||
)
|
||||
|
||||
sched_output = self.scheduler.schedule()
|
||||
|
||||
state = self.scheduler.get_request_state("native-error")
|
||||
assert state is not None
|
||||
assert state.status == DiffusionRequestStatus.FINISHED_ERROR
|
||||
assert state.error == "native allocation bug"
|
||||
assert sched_output.finished_req_ids == {"native-error"}
|
||||
assert sched_output.scheduled_request_ids == []
|
||||
|
||||
def test_diffusion_kv_preemption_retains_allocation(self) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler, num_blocks=3)
|
||||
request = _make_request("preempted")
|
||||
_attach_diffusion_kv(request)
|
||||
self.scheduler.add_request(request)
|
||||
self.scheduler.schedule()
|
||||
manager = self.scheduler._diffusion_kv_manager
|
||||
assert manager is not None
|
||||
metadata = manager.get_metadata("preempted")
|
||||
|
||||
assert self.scheduler.preempt_request("preempted") is True
|
||||
resumed_output = self.scheduler.schedule()
|
||||
|
||||
assert resumed_output.scheduled_new_reqs == []
|
||||
assert _cached_ids(resumed_output) == ["preempted"]
|
||||
assert manager.get_metadata("preempted") == metadata
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status",
|
||||
[
|
||||
DiffusionRequestStatus.FINISHED_COMPLETED,
|
||||
DiffusionRequestStatus.FINISHED_ABORTED,
|
||||
DiffusionRequestStatus.FINISHED_ERROR,
|
||||
],
|
||||
)
|
||||
def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler, num_blocks=3)
|
||||
request = _make_request("terminal")
|
||||
_attach_diffusion_kv(request)
|
||||
self.scheduler.add_request(request)
|
||||
self.scheduler.schedule()
|
||||
manager = self.scheduler._diffusion_kv_manager
|
||||
assert manager is not None
|
||||
assert manager.has_request("terminal") is True
|
||||
|
||||
self.scheduler.finish_requests("terminal", status)
|
||||
|
||||
assert manager.has_request("terminal") is False
|
||||
|
||||
def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler)
|
||||
self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
|
||||
request = _make_request("diffusion-kv")
|
||||
|
||||
with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
|
||||
@@ -656,7 +443,7 @@ class TestRequestScheduler:
|
||||
],
|
||||
)
|
||||
def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
|
||||
_initialize_paged_scheduler(self.scheduler)
|
||||
self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
|
||||
request = _make_request("diffusion-kv")
|
||||
owner = request if owner_name == "request" else request.sampling_params
|
||||
setattr(owner, field_name, object())
|
||||
@@ -1134,45 +921,6 @@ class TestDiffusionEngine:
|
||||
assert engine.scheduler.max_num_running_reqs == 1
|
||||
fake_executor_cls.assert_called_once_with(od_config)
|
||||
|
||||
def test_initializes_paged_scheduler_from_native_control_plane(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
mocker: MockerFixture,
|
||||
) -> None:
|
||||
od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
|
||||
fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
|
||||
scheduler = mocker.Mock()
|
||||
kv_cache_config = object()
|
||||
kv_vllm_config = object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
|
||||
lambda *args, **kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
|
||||
lambda *args, **kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
|
||||
lambda *args, **kwargs: fake_executor_cls,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
|
||||
lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config),
|
||||
)
|
||||
|
||||
engine = DiffusionEngine(od_config, scheduler=scheduler)
|
||||
|
||||
assert engine.scheduler is scheduler
|
||||
scheduler.initialize.assert_called_once_with(
|
||||
od_config,
|
||||
kv_cache_config=kv_cache_config,
|
||||
scheduler_block_size=16,
|
||||
hash_block_size=16,
|
||||
kv_vllm_config=kv_vllm_config,
|
||||
)
|
||||
|
||||
def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
|
||||
scheduler = Scheduler()
|
||||
scheduler.initialize(SimpleNamespace())
|
||||
@@ -1239,21 +987,6 @@ class TestDiffusionEngine:
|
||||
assert output.aborted is True
|
||||
assert output.abort_message == "Request req-finalize aborted."
|
||||
|
||||
def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.scheduler = StepScheduler()
|
||||
engine.scheduler.initialize(SimpleNamespace())
|
||||
|
||||
req_id = engine.scheduler.add_request(_make_request("req-error"))
|
||||
engine.scheduler._finish_requests(
|
||||
{req_id: DiffusionRequestStatus.FINISHED_ERROR},
|
||||
{req_id: "KV request cannot fit"},
|
||||
)
|
||||
|
||||
output = engine._finalize_finished_request(req_id)
|
||||
|
||||
assert output.error == "KV request cannot fit"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
|
||||
@@ -13,7 +13,11 @@ import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import (
|
||||
DiffusionWorker,
|
||||
_create_diffusion_worker_vllm_config,
|
||||
_make_diffusion_vllm_model_config,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
|
||||
|
||||
@@ -54,6 +58,34 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
|
||||
return worker
|
||||
|
||||
|
||||
def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
model="dummy",
|
||||
dtype=torch.bfloat16,
|
||||
quantization_config=build_quant_config(
|
||||
{
|
||||
"quant_method": "modelopt",
|
||||
"quant_algo": "FP8",
|
||||
"ignore": [],
|
||||
}
|
||||
),
|
||||
tf_model_config=SimpleNamespace(),
|
||||
enforce_eager=True,
|
||||
is_moe=False,
|
||||
)
|
||||
|
||||
model_config = _make_diffusion_vllm_model_config(od_config)
|
||||
|
||||
assert model_config.dtype is torch.bfloat16
|
||||
assert model_config.quantization == "modelopt"
|
||||
assert model_config.quantization_config is od_config.quantization_config
|
||||
assert model_config.is_quantized()
|
||||
|
||||
|
||||
class TestDiffusionWorkerSleep:
|
||||
"""Test DiffusionWorker.sleep method."""
|
||||
|
||||
@@ -297,3 +329,49 @@ class TestDiffusionWorkerWakeUp:
|
||||
mock_buffer2.data.copy_.assert_not_called()
|
||||
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestWorkerVllmConfigAdditionalConfig:
|
||||
"""Test worker-side VllmConfig construction with additional_config."""
|
||||
|
||||
def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture):
|
||||
mock_vllm_config = mocker.Mock()
|
||||
mock_vllm_cls = mocker.patch(
|
||||
"vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
|
||||
return_value=mock_vllm_config,
|
||||
)
|
||||
od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}})
|
||||
|
||||
result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
|
||||
|
||||
assert result is mock_vllm_config
|
||||
assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config
|
||||
|
||||
def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config(
|
||||
self, mocker: MockerFixture
|
||||
):
|
||||
mock_vllm_config = mocker.Mock()
|
||||
mock_vllm_cls = mocker.patch(
|
||||
"vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
|
||||
side_effect=[
|
||||
TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"),
|
||||
mock_vllm_config,
|
||||
],
|
||||
)
|
||||
od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
|
||||
|
||||
result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
|
||||
|
||||
assert result is mock_vllm_config
|
||||
assert mock_vllm_cls.call_count == 2
|
||||
assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config
|
||||
|
||||
def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture):
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
|
||||
side_effect=TypeError("additional_config values must be hashable"),
|
||||
)
|
||||
od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
|
||||
|
||||
with pytest.raises(TypeError, match="hashable"):
|
||||
_create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
|
||||
@@ -413,26 +413,6 @@ class TestRequestModeDispatch:
|
||||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||||
executor.collective_rpc.assert_called_once()
|
||||
|
||||
def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(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.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||||
scheduler_output = _make_sched_output("A", "B")
|
||||
for new_req in scheduler_output.scheduled_new_reqs:
|
||||
new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
|
||||
|
||||
result = executor.execute_request(scheduler_output)
|
||||
|
||||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||||
forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
|
||||
assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
|
||||
assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
|
||||
@@ -1,159 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.advanced_model,
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.gpu,
|
||||
]
|
||||
|
||||
|
||||
def _make_sana_wm_e2e_image() -> Any:
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.sana_wm import SANA_WM_OUTPUT_HEIGHT, SANA_WM_OUTPUT_WIDTH
|
||||
|
||||
return Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
|
||||
|
||||
|
||||
def _sana_wm_e2e_tensor_parallel_size() -> int:
|
||||
return int(
|
||||
os.environ.get(
|
||||
"SANA_WM_E2E_TENSOR_PARALLEL_SIZE",
|
||||
os.environ.get("SANA_WM_E2E_TP", "1"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _coerce_video_array(video: Any) -> np.ndarray:
|
||||
import torch
|
||||
|
||||
if isinstance(video, list):
|
||||
assert video, "SANA-WM e2e produced an empty video list."
|
||||
video = video[0]
|
||||
if isinstance(video, torch.Tensor):
|
||||
video = video.detach().cpu().float().numpy()
|
||||
video_array = np.asarray(video)
|
||||
if video_array.ndim == 5:
|
||||
assert video_array.shape[0] == 1
|
||||
video_array = video_array[0]
|
||||
assert video_array.ndim == 4
|
||||
return video_array
|
||||
|
||||
|
||||
def _assert_sana_wm_e2e_shape(
|
||||
video: np.ndarray,
|
||||
*,
|
||||
output_type: str,
|
||||
num_frames: int,
|
||||
) -> None:
|
||||
if output_type == "latent":
|
||||
# Latent output is channel-first: (C, T, H, W) after removing the
|
||||
# batch dimension. Stage-1 emits LTX-2 latents with 128 channels.
|
||||
assert video.shape[0] == 128
|
||||
assert 0 < video.shape[1] <= num_frames
|
||||
return
|
||||
assert 0 < video.shape[0] <= num_frames
|
||||
|
||||
|
||||
def _run_sana_wm_e2e(*, output_type: str) -> np.ndarray:
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.sana_wm import (
|
||||
SANA_WM_MODEL_ID,
|
||||
SANA_WM_OUTPUT_HEIGHT,
|
||||
SANA_WM_OUTPUT_WIDTH,
|
||||
)
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("Sana-WM e2e requires CUDA.")
|
||||
|
||||
model = os.environ.get("SANA_WM_E2E_MODEL", SANA_WM_MODEL_ID)
|
||||
model_class_name = os.environ.get("SANA_WM_E2E_MODEL_CLASS", "SanaWmPipeline")
|
||||
num_frames = int(os.environ.get("SANA_WM_E2E_NUM_FRAMES", "9"))
|
||||
stage1_steps = int(
|
||||
os.environ.get(
|
||||
"SANA_WM_E2E_STAGE1_STEPS",
|
||||
os.environ.get("SANA_WM_E2E_NUM_INFERENCE_STEPS", "1"),
|
||||
)
|
||||
)
|
||||
action = os.environ.get("SANA_WM_E2E_ACTION", f"w-{max(num_frames - 1, 1)}")
|
||||
tensor_parallel_size = _sana_wm_e2e_tensor_parallel_size()
|
||||
image = _make_sana_wm_e2e_image()
|
||||
omni = Omni(
|
||||
model=model,
|
||||
model_class_name=model_class_name,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
extra_args: dict[str, Any] = {}
|
||||
# The native Stage-1 path is latent-token capped; raise it for full-size
|
||||
# e2e runs (704x1280 exceeds the default 4096-token cap).
|
||||
native_max_tokens = os.environ.get("SANA_WM_E2E_NATIVE_MAX_TOKENS", "").strip()
|
||||
if native_max_tokens:
|
||||
extra_args["sana_wm_native_max_tokens"] = int(native_max_tokens)
|
||||
# Return the requested output type so the shape assertion matches the
|
||||
# pipeline output instead of the raw Stage-1 latent default.
|
||||
|
||||
output = omni.generate(
|
||||
{
|
||||
"prompt": "A slow forward camera move through a quiet city street.",
|
||||
"multi_modal_data": {"image": image},
|
||||
"sana_wm": {
|
||||
"action": action,
|
||||
"num_frames": num_frames,
|
||||
"translation_speed": 0.055,
|
||||
"rotation_speed_deg": 1.2,
|
||||
# Avoid the optional Pi3X dependency by passing deterministic
|
||||
# camera intrinsics directly.
|
||||
"intrinsics": {
|
||||
"fx": SANA_WM_OUTPUT_WIDTH / 2,
|
||||
"fy": SANA_WM_OUTPUT_WIDTH / 2,
|
||||
"cx": SANA_WM_OUTPUT_WIDTH / 2,
|
||||
"cy": SANA_WM_OUTPUT_HEIGHT / 2,
|
||||
},
|
||||
},
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=SANA_WM_OUTPUT_HEIGHT,
|
||||
width=SANA_WM_OUTPUT_WIDTH,
|
||||
num_frames=num_frames,
|
||||
seed=0,
|
||||
fps=16,
|
||||
num_inference_steps=stage1_steps,
|
||||
guidance_scale=1.0,
|
||||
guidance_scale_provided=True,
|
||||
output_type=output_type,
|
||||
extra_args=extra_args,
|
||||
),
|
||||
)
|
||||
|
||||
request_output = output[0] if isinstance(output, list) else output
|
||||
assert isinstance(request_output, OmniRequestOutput)
|
||||
assert request_output.error is None
|
||||
assert request_output.images
|
||||
|
||||
frames = request_output.images[0]
|
||||
video = _coerce_video_array(frames)
|
||||
_assert_sana_wm_e2e_shape(video, output_type=output_type, num_frames=num_frames)
|
||||
return video
|
||||
|
||||
|
||||
def test_sana_wm_native_generates_video() -> None:
|
||||
output_type = os.environ.get("SANA_WM_E2E_OUTPUT_TYPE", "np")
|
||||
video = _run_sana_wm_e2e(output_type=output_type)
|
||||
if output_type == "latent":
|
||||
assert video.ndim == 4
|
||||
else:
|
||||
assert video.shape[-1] in (3, 4)
|
||||
@@ -1,112 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E online tests for MiniMax Music 3 text-to-music.
|
||||
|
||||
This is a music model on the speech endpoint, so the request shape differs
|
||||
from every other TTS model here: ``input`` carries the lyrics, ``instructions``
|
||||
carries the caption that decides genre and arrangement, and both are required.
|
||||
|
||||
Weekly rather than per-PR: the checkpoint is 27 GB across five components and a
|
||||
single request occupies two decode rows, because classifier-free guidance is
|
||||
mandatory for this model. Rejection cases live in
|
||||
``tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py``.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
MODEL = "MiniMaxAI/MiniMax-Music3"
|
||||
DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 600.0
|
||||
|
||||
# 25 frames per second, so 250 frames is ten seconds and 750 is thirty. The
|
||||
# byte floors sit below the true payload (32 kHz stereo, 16-bit) to leave room
|
||||
# for the model ending a song before the cap.
|
||||
_FRAMES_10S = 250
|
||||
_FRAMES_30S = 750
|
||||
_MIN_BYTES_10S = 500_000
|
||||
_MIN_BYTES_30S = 1_500_000
|
||||
|
||||
LYRICS = (
|
||||
"[Verse]\nWalking down the empty street at midnight\n"
|
||||
"Streetlights flicker like a broken dream\n"
|
||||
"[Chorus]\nAnd I keep on walking\nTill the morning finds me"
|
||||
)
|
||||
|
||||
CAPTION = (
|
||||
"A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, "
|
||||
"soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. "
|
||||
"Intimate bedroom production, gentle tape saturation, no bright cymbals."
|
||||
)
|
||||
|
||||
|
||||
tts_server_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
|
||||
server_args=["--trust-remote-code", "--disable-log-stats"],
|
||||
),
|
||||
id="minimax_music3",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_to_music_001(omni_server, openai_client) -> None:
|
||||
"""Baseline: lyrics plus caption in, one ten-second song out.
|
||||
|
||||
Deploy Setting: minimax_music3.yaml
|
||||
Input Modal: text (lyrics) + text (caption)
|
||||
Output Modal: audio, 32 kHz stereo
|
||||
Input Setting: stream=False, seed pinned, 250 frames
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": LYRICS,
|
||||
"instructions": CAPTION,
|
||||
"seed": 1,
|
||||
"max_new_tokens": _FRAMES_10S,
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_BYTES_10S,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_to_music_multi_window_002(omni_server, openai_client) -> None:
|
||||
"""A song long enough to span several acoustic windows.
|
||||
|
||||
The acoustic stage decodes overlapping 200-frame windows on a 100-frame
|
||||
hop and crops the overlap back off, so a request that fits in one window
|
||||
never exercises the stitching. 750 frames covers seven of them.
|
||||
|
||||
Deploy Setting: minimax_music3.yaml
|
||||
Input Setting: stream=False, seed pinned, 750 frames
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": LYRICS,
|
||||
"instructions": CAPTION,
|
||||
"seed": 1,
|
||||
"max_new_tokens": _FRAMES_30S,
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_BYTES_30S,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
@@ -1,106 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``).
|
||||
|
||||
Boots the server straight from the model id — the Stage-1 repo's
|
||||
``model_index.json`` names ``SanaWmPipeline``, so the class resolves on its own
|
||||
— submits one async ``/v1/videos`` job with a first-frame reference image, and
|
||||
asserts the job completes and returns video bytes.
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_sana_wm.py -m "advanced_model and diffusion" --run-level=advanced_model
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from vllm_omni.diffusion.models.sana_wm import (
|
||||
SANA_WM_MODEL_ID,
|
||||
SANA_WM_OUTPUT_HEIGHT,
|
||||
SANA_WM_OUTPUT_WIDTH,
|
||||
)
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
MODEL = SANA_WM_MODEL_ID
|
||||
PROMPT = "A slow forward camera move through a quiet city street."
|
||||
NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
|
||||
|
||||
# Smoke knobs: a handful of frames / steps so the job finishes quickly.
|
||||
SMOKE_NUM_FRAMES = 9
|
||||
SMOKE_NUM_INFERENCE_STEPS = 2
|
||||
# First-frame I2V needs a camera trajectory; "w-<n>" is the forward-move action DSL.
|
||||
# Passed via the ``sana_wm`` form field (a JSON string parsed server-side); explicit
|
||||
# intrinsics avoid the optional Pi3X camera-calibration dependency.
|
||||
SANA_WM_PARAMS = {
|
||||
"action": f"w-{SMOKE_NUM_FRAMES - 1}",
|
||||
"translation_speed": 0.055,
|
||||
"rotation_speed_deg": 1.2,
|
||||
"intrinsics": {
|
||||
"fx": SANA_WM_OUTPUT_WIDTH / 2,
|
||||
"fy": SANA_WM_OUTPUT_WIDTH / 2,
|
||||
"cx": SANA_WM_OUTPUT_WIDTH / 2,
|
||||
"cy": SANA_WM_OUTPUT_HEIGHT / 2,
|
||||
},
|
||||
}
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _first_frame_data_url() -> str:
|
||||
"""A solid-color first frame encoded as a PNG ``data:`` URL for I2V."""
|
||||
from PIL import Image
|
||||
|
||||
image = Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
return f"data:image/png;base64,{encoded}"
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
"""Single default server row."""
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(model=model),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_image_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""Default SANA-WM I2V smoke: async ``/v1/videos`` job completes and returns video bytes."""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"image_reference": _first_frame_data_url(),
|
||||
"form_data": {
|
||||
"prompt": PROMPT,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
# The camera payload rides inside extra_params; there is no
|
||||
# top-level sana_wm form field on the video endpoints, so sending
|
||||
# one is silently dropped and the request fails in the
|
||||
# preprocessor for missing camera control.
|
||||
"extra_params": json.dumps({"sana_wm": SANA_WM_PARAMS}),
|
||||
"height": SANA_WM_OUTPUT_HEIGHT,
|
||||
"width": SANA_WM_OUTPUT_WIDTH,
|
||||
"num_frames": SMOKE_NUM_FRAMES,
|
||||
"fps": 16,
|
||||
"num_inference_steps": SMOKE_NUM_INFERENCE_STEPS,
|
||||
"guidance_scale": 5.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_video_diffusion_request(request_config)
|
||||
@@ -371,106 +371,3 @@ def test_build_add_request_message_releases_preselected_replica_on_preprocess_er
|
||||
)
|
||||
|
||||
assert stage_pool.get_bound_replica_id("req-error") is None
|
||||
|
||||
|
||||
def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture):
|
||||
"""A guided request is all-or-nothing.
|
||||
|
||||
Admitting the parent before its companion is built leaves an orphan: a
|
||||
model whose guidance is mandatory never completes the pair, so the request
|
||||
holds scheduler and KV capacity for the scheduler's whole hold budget and
|
||||
then produces no audio. The build has to raise before anything is enqueued.
|
||||
"""
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin"))
|
||||
engine.supported_tasks = ("generate",)
|
||||
engine.input_processor = mocker.Mock()
|
||||
engine.request_queue = mocker.Mock()
|
||||
|
||||
with pytest.raises(ValueError, match="null twin"):
|
||||
engine._build_cfg_companions(
|
||||
parent_id="req",
|
||||
original_prompt={"prompt": "positive"},
|
||||
stage0_params=params,
|
||||
sampling_params_list=[params],
|
||||
)
|
||||
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
|
||||
|
||||
def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture):
|
||||
"""The same holds when expansion succeeds but input processing does not."""
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
engine.prompt_expand_func = lambda *_args: [
|
||||
ExpandedPrompt(
|
||||
prompt={"prompt": "negative"},
|
||||
role="cfg_text",
|
||||
request_id_suffix="__cfg_text",
|
||||
)
|
||||
]
|
||||
engine.supported_tasks = ("generate",)
|
||||
engine.input_processor = mocker.Mock()
|
||||
engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt")
|
||||
engine.request_queue = mocker.Mock()
|
||||
|
||||
with pytest.raises(RuntimeError, match="tokenizer rejected"):
|
||||
engine._build_cfg_companions(
|
||||
parent_id="req",
|
||||
original_prompt={"prompt": "positive"},
|
||||
stage0_params=params,
|
||||
sampling_params_list=[params],
|
||||
)
|
||||
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
|
||||
|
||||
def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture):
|
||||
"""Construction and admission are separate steps."""
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
engine.prompt_expand_func = lambda *_args: [
|
||||
ExpandedPrompt(
|
||||
prompt={"prompt": "negative"},
|
||||
role="cfg_text",
|
||||
request_id_suffix="__cfg_text",
|
||||
)
|
||||
]
|
||||
engine.supported_tasks = ("generate",)
|
||||
engine.input_processor = mocker.Mock()
|
||||
engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
|
||||
engine.request_queue = mocker.Mock()
|
||||
|
||||
companions = engine._build_cfg_companions(
|
||||
parent_id="req",
|
||||
original_prompt={"prompt": "positive"},
|
||||
stage0_params=params,
|
||||
sampling_params_list=[params],
|
||||
)
|
||||
|
||||
assert len(companions) == 1
|
||||
assert companions[0].companion_id == "req__cfg_text"
|
||||
assert companions[0].parent_id == "req"
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
|
||||
|
||||
def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture):
|
||||
"""An unguided request expands to no companions and admits normally."""
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
engine.prompt_expand_func = lambda *_args: []
|
||||
engine.supported_tasks = ("generate",)
|
||||
engine.input_processor = mocker.Mock()
|
||||
engine.request_queue = mocker.Mock()
|
||||
|
||||
assert (
|
||||
engine._build_cfg_companions(
|
||||
parent_id="req",
|
||||
original_prompt={"prompt": "positive"},
|
||||
stage0_params=params,
|
||||
sampling_params_list=[params],
|
||||
)
|
||||
== []
|
||||
)
|
||||
engine.request_queue.sync_q.put.assert_not_called()
|
||||
@@ -8,23 +8,8 @@ from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic.fields import FieldInfo
|
||||
from vllm.config import CacheConfig as VllmCacheConfig
|
||||
from vllm.config import CompilationConfig as VllmCompilationConfig
|
||||
from vllm.config import LoadConfig as VllmLoadConfig
|
||||
from vllm.config import ParallelConfig as VllmParallelConfig
|
||||
from vllm.config import ProfilerConfig as VllmProfilerConfig
|
||||
from vllm.config import SchedulerConfig as VllmSchedulerConfig
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
|
||||
from vllm_omni.config.omni_config import (
|
||||
_LLM_STAGE_ENGINE_FIELDS,
|
||||
OmniStageCacheConfig,
|
||||
OmniStageLoadConfig,
|
||||
OmniStageParallelConfig,
|
||||
OmniStageSchedulerConfig,
|
||||
VllmOmniARStageConfig,
|
||||
VllmOmniConfig,
|
||||
)
|
||||
from vllm_omni.config.omni_config import VllmOmniConfig
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
|
||||
from vllm_omni.config.stage_config import (
|
||||
DeployConfig,
|
||||
@@ -68,30 +53,6 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str,
|
||||
|
||||
_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
|
||||
_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
|
||||
_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
|
||||
{
|
||||
"active_stream_window",
|
||||
"codec_frame_rate_hz",
|
||||
"custom_voice_dir",
|
||||
"devices",
|
||||
"disable_autocast",
|
||||
"enable_multithread_weight_load",
|
||||
"env",
|
||||
"has_sampling_extra_args",
|
||||
"log_level",
|
||||
"log_stats",
|
||||
"model_arch",
|
||||
"model_subdir",
|
||||
"num_gpus",
|
||||
"num_replicas",
|
||||
"num_weight_load_threads",
|
||||
"omni_kv_config",
|
||||
"parallel_config",
|
||||
"subtalker_sampling_params",
|
||||
"task_type",
|
||||
"tokenizer_subdir",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@@ -236,100 +197,6 @@ def _legacy_and_typed_stages(
|
||||
return legacy_stages, omni_config
|
||||
|
||||
|
||||
def test_llm_stage_engine_field_schema_tracks_upstream_engine_args():
|
||||
upstream_engine_fields = frozenset(field.name for field in fields(EngineArgs))
|
||||
|
||||
assert _LLM_STAGE_ENGINE_FIELDS - upstream_engine_fields == _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS
|
||||
|
||||
|
||||
def test_typed_llm_projection_discovers_explicit_upstream_config_fields():
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
|
||||
load_config=OmniStageLoadConfig(ignore_patterns=["*.bin"]),
|
||||
cache_config=OmniStageCacheConfig(block_size=32, cache_dtype="fp8"),
|
||||
scheduler_config=OmniStageSchedulerConfig(
|
||||
prefill_schedule_interval=2,
|
||||
policy="priority",
|
||||
),
|
||||
parallel_config=OmniStageParallelConfig(
|
||||
data_parallel_master_ip="10.0.0.1",
|
||||
enable_dbo=True,
|
||||
),
|
||||
)
|
||||
|
||||
engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
|
||||
|
||||
assert engine_args["ignore_patterns"] == ["*.bin"]
|
||||
assert engine_args["block_size"] == 32
|
||||
assert engine_args["kv_cache_dtype"] == "fp8"
|
||||
assert engine_args["prefill_schedule_interval"] == 2
|
||||
assert engine_args["scheduling_policy"] == "priority"
|
||||
assert engine_args["data_parallel_address"] == "10.0.0.1"
|
||||
assert engine_args["enable_dbo"] is True
|
||||
|
||||
|
||||
def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
|
||||
)
|
||||
|
||||
engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
|
||||
|
||||
inherited_defaults = {
|
||||
"ignore_patterns",
|
||||
"block_size",
|
||||
"kv_cache_dtype",
|
||||
"prefill_schedule_interval",
|
||||
"scheduling_policy",
|
||||
"data_parallel_address",
|
||||
"data_parallel_rank",
|
||||
"data_parallel_size_local",
|
||||
"data_parallel_rpc_port",
|
||||
"enable_dbo",
|
||||
}
|
||||
assert inherited_defaults.isdisjoint(engine_args)
|
||||
|
||||
|
||||
def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary():
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
|
||||
scheduler_config=OmniStageSchedulerConfig(scheduler_cls="test.CustomScheduler"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"no EngineArgs projection: scheduler_cls"):
|
||||
stage_init_utils._project_omni_stage_engine_args(stage_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_id", [0, 1], ids=["ar", "generation"])
|
||||
def test_typed_llm_engine_args_preserve_upstream_config_objects(tmp_path, stage_id):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
compilation_config = VllmCompilationConfig(backend="eager")
|
||||
profiler_config = VllmProfilerConfig(profiler="cuda")
|
||||
stage_prefix = f"stage_{stage_id}_"
|
||||
omni_config = VllmOmniConfig.from_pipeline_config(
|
||||
pipeline,
|
||||
user_deploy_config=copy.deepcopy(deploy),
|
||||
cli_overrides={
|
||||
"model": model,
|
||||
f"{stage_prefix}compilation_config": compilation_config,
|
||||
f"{stage_prefix}profiler_config": profiler_config,
|
||||
},
|
||||
)
|
||||
|
||||
stage_config = omni_config.stage_by_id(stage_id)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(stage_config, model)
|
||||
|
||||
assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
|
||||
assert isinstance(typed_args["compilation_config"], VllmCompilationConfig)
|
||||
assert typed_args["compilation_config"] is not stage_config.compilation_config
|
||||
assert typed_args["compilation_config"].backend == "eager"
|
||||
|
||||
assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
|
||||
assert isinstance(typed_args["profiler_config"], VllmProfilerConfig)
|
||||
assert typed_args["profiler_config"] is not stage_config.profiler_config
|
||||
assert typed_args["profiler_config"].profiler == "cuda"
|
||||
|
||||
|
||||
def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
@@ -355,13 +222,6 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
|
||||
assert {name: typed_args[name] for name in legacy_args} == legacy_args
|
||||
|
||||
thinker_args = typed_args_by_stage[0]
|
||||
inherited_vllm_fields = {
|
||||
field.name
|
||||
for config_cls in (VllmLoadConfig, VllmCacheConfig, VllmSchedulerConfig, VllmParallelConfig)
|
||||
for field in fields(config_cls)
|
||||
}
|
||||
topology_projected_fields = {"scheduler_cls"}
|
||||
assert (inherited_vllm_fields - _LLM_STAGE_ENGINE_FIELDS - topology_projected_fields).isdisjoint(thinker_args)
|
||||
assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
|
||||
assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
|
||||
assert thinker_args["stage_id"] == 0
|
||||
|
||||
@@ -1,810 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""API-server surface guards for the OpenAI Omni entrypoint.
|
||||
|
||||
Purpose
|
||||
-------
|
||||
Freeze the *observable* server contract so helper moves (P0.2), route
|
||||
extraction (P0.3), and later endpoint-family PRs cannot silently change:
|
||||
|
||||
* which routes exist,
|
||||
* how upstream routes are overridden at assembly time,
|
||||
* unavailable / missing-handler envelopes,
|
||||
* shared engine-error JSON fields,
|
||||
* app-state keys wired by ``omni_init_app_state``.
|
||||
|
||||
CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON.
|
||||
|
||||
Common wrong changes these tests are meant to catch
|
||||
----------------------------------------------------
|
||||
* Route accidentally removed, renamed, method-changed, or double-registered
|
||||
after a move (census / assembled-app checks).
|
||||
* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
|
||||
* Assembly forgets to strip upstream ``/v1/chat/completions`` (batch) /
|
||||
``/v1/models``, or also deletes unrelated upstream routes.
|
||||
* Profiler / Omni router not mounted; storage not started; Omni engine
|
||||
exception handlers not registered.
|
||||
* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
|
||||
* Missing-handler WS/HTTP responses change status, JSON shape, or wording
|
||||
(clients treat these as capability probes).
|
||||
* Engine error payloads lose ``request_id`` / ``error_stage_id``.
|
||||
* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a
|
||||
handler as ``None`` while the route is still registered.
|
||||
|
||||
Not covered here (add elsewhere if needed)
|
||||
------------------------------------------
|
||||
* Deep per-endpoint success behavior (family UTs).
|
||||
* Byte-identical helper bodies after a move (P0.2 AST integrity).
|
||||
* Live GPU / speech byte-identity checks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.routing import APIRoute, APIWebSocketRoute
|
||||
from fastapi.testclient import TestClient
|
||||
from starlette.datastructures import State
|
||||
from starlette.requests import Request
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
|
||||
from vllm_omni.entrypoints.openai import api_server
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
|
||||
# would invent ("HEAD", path) entries for every GET and fail the current census.
|
||||
# TODO: extend this set and the expected census if @router.head / @router.options appear.
|
||||
_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
|
||||
|
||||
# Exact Omni router census on main. Update intentionally when routes change.
|
||||
_EXPECTED_ROUTER_ROUTES = {
|
||||
("POST", "/v1/chat/completions"),
|
||||
("POST", "/v1/chat/completions/batch"),
|
||||
("POST", "/v1/audio/speech"),
|
||||
("POST", "/v1/audio/speech/batch"),
|
||||
("POST", "/v1/audio/generate"),
|
||||
("GET", "/v1/audio/voices"),
|
||||
("POST", "/v1/audio/voices"),
|
||||
("DELETE", "/v1/audio/voices/{name}"),
|
||||
("WEBSOCKET", "/v1/audio/speech/stream"),
|
||||
("WEBSOCKET", "/v1/video/chat/stream"),
|
||||
("WEBSOCKET", "/v1/realtime/video"),
|
||||
("WEBSOCKET", "/v1/realtime"),
|
||||
("WEBSOCKET", "/v1/realtime/robot/openpi"),
|
||||
("WEBSOCKET", "/v1/duplex"),
|
||||
("GET", "/health"),
|
||||
("GET", "/v1/models"),
|
||||
("POST", "/v1/images/generations"),
|
||||
("POST", "/v1/images/edits"),
|
||||
("POST", "/v1/videos"),
|
||||
("POST", "/v1/videos/sync"),
|
||||
("GET", "/v1/videos"),
|
||||
("GET", "/v1/videos/{video_id}"),
|
||||
("DELETE", "/v1/videos/{video_id}"),
|
||||
("GET", "/v1/videos/{video_id}/content"),
|
||||
("POST", "/v1/omni/sleep"),
|
||||
("POST", "/v1/omni/wakeup"),
|
||||
}
|
||||
_EXPECTED_PROFILER_ROUTES = {
|
||||
("POST", "/start_profile"),
|
||||
("POST", "/stop_profile"),
|
||||
}
|
||||
|
||||
# App-state keys that must remain available across refactor phases.
|
||||
# Presence is not enough: "not wired" is written as ``state.X = None``.
|
||||
_DIFFUSION_APP_STATE_KEYS = {
|
||||
"engine_client",
|
||||
"log_stats",
|
||||
"args",
|
||||
"sleeping_stages",
|
||||
"stage_configs",
|
||||
"vllm_config",
|
||||
"diffusion_engine",
|
||||
"openai_serving_models",
|
||||
"serving_tokenization",
|
||||
"openai_serving_chat",
|
||||
"openai_serving_chat_batch",
|
||||
"openai_serving_speech",
|
||||
"openai_serving_audio_generate",
|
||||
"openai_serving_video",
|
||||
"openai_streaming_video_output",
|
||||
"openai_serving_duplex",
|
||||
"openai_streaming_speech",
|
||||
"openai_streaming_video",
|
||||
"openai_serving_realtime_robot",
|
||||
"enable_server_load_tracking",
|
||||
"server_load_metrics",
|
||||
}
|
||||
_DIFFUSION_MUST_BE_NONE = {
|
||||
"vllm_config",
|
||||
"serving_tokenization",
|
||||
"openai_serving_duplex",
|
||||
"openai_streaming_speech",
|
||||
"openai_streaming_video",
|
||||
}
|
||||
_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE
|
||||
_MULTISTAGE_APP_STATE_KEYS = {
|
||||
"engine_client",
|
||||
"log_stats",
|
||||
"args",
|
||||
"sleeping_stages",
|
||||
"stage_configs",
|
||||
"vllm_config",
|
||||
"openai_serving_models",
|
||||
"serving_tokenization",
|
||||
"online_renderer",
|
||||
"openai_serving_chat",
|
||||
"openai_serving_chat_batch",
|
||||
"openai_serving_speech",
|
||||
"openai_serving_audio_generate",
|
||||
"openai_streaming_speech",
|
||||
"openai_streaming_video",
|
||||
"openai_serving_duplex",
|
||||
"openai_serving_realtime",
|
||||
"openai_serving_video",
|
||||
"openai_serving_realtime_robot",
|
||||
"enable_server_load_tracking",
|
||||
"server_load_metrics",
|
||||
}
|
||||
_MULTISTAGE_MUST_BE_NONE = {
|
||||
"openai_serving_duplex",
|
||||
"openai_serving_realtime_robot",
|
||||
}
|
||||
_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
|
||||
|
||||
|
||||
def _route_entries(routes) -> list[tuple[str, str]]:
|
||||
entries: list[tuple[str, str]] = []
|
||||
for route in routes:
|
||||
if isinstance(route, APIRoute):
|
||||
for method in sorted((route.methods or set()) & _HTTP_METHODS):
|
||||
entries.append((method, route.path))
|
||||
elif isinstance(route, APIWebSocketRoute):
|
||||
entries.append(("WEBSOCKET", route.path))
|
||||
return entries
|
||||
|
||||
|
||||
def _route_manifest(routes) -> set[tuple[str, str]]:
|
||||
return set(_route_entries(routes))
|
||||
|
||||
|
||||
def _assert_unique_route_entries(routes) -> None:
|
||||
entries = _route_entries(routes)
|
||||
duplicates = sorted({item for item in entries if entries.count(item) > 1})
|
||||
assert not duplicates, f"duplicate route registrations: {duplicates}"
|
||||
|
||||
|
||||
def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]:
|
||||
return [
|
||||
route
|
||||
for route in routes
|
||||
if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set())
|
||||
]
|
||||
|
||||
|
||||
def _single_http_route(routes, method: str, path: str) -> APIRoute:
|
||||
matches = _matching_http_routes(routes, method, path)
|
||||
assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}"
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _assert_app_state_snapshot(
|
||||
state,
|
||||
*,
|
||||
expected_keys: set[str],
|
||||
must_be_wired: set[str],
|
||||
must_be_none: set[str],
|
||||
) -> None:
|
||||
assert must_be_wired.isdisjoint(must_be_none)
|
||||
assert must_be_wired | must_be_none <= expected_keys
|
||||
present = {key for key in expected_keys if hasattr(state, key)}
|
||||
assert present == expected_keys
|
||||
not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None)
|
||||
assert not not_wired, f"app.state keys registered but not wired: {not_wired}"
|
||||
unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None)
|
||||
assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}"
|
||||
|
||||
|
||||
def _minimal_args(**overrides) -> SimpleNamespace:
|
||||
args = SimpleNamespace(
|
||||
tool_parser_plugin=None,
|
||||
reasoning_parser_plugin=None,
|
||||
reasoning_parser=None,
|
||||
structured_outputs_config=SimpleNamespace(reasoning_parser=None),
|
||||
enable_ssl_refresh=False,
|
||||
host="127.0.0.1",
|
||||
port=8000,
|
||||
uvicorn_log_level="info",
|
||||
disable_uvicorn_access_log=True,
|
||||
ssl_keyfile=None,
|
||||
ssl_certfile=None,
|
||||
ssl_ca_certs=None,
|
||||
ssl_cert_reqs=0,
|
||||
ssl_ciphers=None,
|
||||
h11_max_incomplete_event_size=None,
|
||||
h11_max_header_count=None,
|
||||
log_config_file=None,
|
||||
disable_access_log_for_endpoints=None,
|
||||
disable_log_stats=True,
|
||||
enable_log_requests=False,
|
||||
max_log_len=None,
|
||||
served_model_name=None,
|
||||
model="demo-model",
|
||||
enable_server_load_tracking=False,
|
||||
log_error_stack=False,
|
||||
chat_template=None,
|
||||
chat_template_content_format="auto",
|
||||
trust_request_chat_template=False,
|
||||
tool_server=None,
|
||||
lora_modules=None,
|
||||
response_role="assistant",
|
||||
default_chat_template_kwargs=None,
|
||||
return_tokens_as_token_ids=False,
|
||||
enable_auto_tool_choice=False,
|
||||
exclude_tools_when_tool_choice_none=False,
|
||||
tool_call_parser=None,
|
||||
enable_prompt_tokens_details=False,
|
||||
enable_force_include_usage=False,
|
||||
enable_log_outputs=False,
|
||||
enable_log_deltas=False,
|
||||
tokens_only=False,
|
||||
shutdown_timeout=0,
|
||||
stage_configs_path=None,
|
||||
deploy_config=None,
|
||||
)
|
||||
for key, value in overrides.items():
|
||||
setattr(args, key, value)
|
||||
return args
|
||||
|
||||
|
||||
def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request:
|
||||
return Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"path": path,
|
||||
"headers": [],
|
||||
"app": app,
|
||||
"state": {},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _FakeSocket:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeEngineClient:
|
||||
def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None:
|
||||
self.stage_configs = stage_configs if stage_configs is not None else []
|
||||
self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {}
|
||||
self.model_config = SimpleNamespace()
|
||||
self.input_processor = object()
|
||||
self.renderer = object()
|
||||
self.errored = False
|
||||
self.vllm_config = vllm_config
|
||||
|
||||
async def get_vllm_config(self):
|
||||
return self.vllm_config
|
||||
|
||||
async def get_supported_tasks(self) -> tuple[str, ...]:
|
||||
return ("generate",)
|
||||
|
||||
async def get_tokenizer(self):
|
||||
return SimpleNamespace(chat_template="dummy")
|
||||
|
||||
async def check_health(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _marker(name: str):
|
||||
return SimpleNamespace(name=name)
|
||||
|
||||
|
||||
def test_router_route_manifest_is_explicit_and_complete() -> None:
|
||||
"""Lock the Omni + profiler router census.
|
||||
|
||||
Fails if a route is added/removed/renamed, a method changes, or a
|
||||
WebSocket path disappears during helper/router moves. Also fails if the
|
||||
same public route is registered more than once, which a set manifest alone
|
||||
would collapse.
|
||||
"""
|
||||
# TODO(P0.3): retarget this object-level census when routes move onto family
|
||||
# routers. The assembled-app census is the client-facing lock that must survive.
|
||||
_assert_unique_route_entries(api_server.router.routes)
|
||||
_assert_unique_route_entries(api_server.profiler_router.routes)
|
||||
assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES
|
||||
assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES
|
||||
|
||||
|
||||
def test_router_openapi_paths_cover_http_manifest() -> None:
|
||||
"""Lock that OpenAPI advertises every HTTP census route/method.
|
||||
|
||||
Fails if a decorator still registers a route but OpenAPI loses the path
|
||||
(or method), e.g. after moving route bodies / response models.
|
||||
"""
|
||||
app = FastAPI()
|
||||
app.include_router(api_server.router)
|
||||
app.include_router(api_server.profiler_router)
|
||||
paths = app.openapi().get("paths", {})
|
||||
|
||||
http_routes = {
|
||||
(method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET"
|
||||
}
|
||||
for method, path in sorted(http_routes):
|
||||
assert path in paths, f"missing OpenAPI path {path}"
|
||||
assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
|
||||
"""Lock worker assembly: override, mount, storage, handlers, full census.
|
||||
|
||||
Fails if assembly stops replacing upstream chat/batch/models, deletes
|
||||
unrelated upstream routes, forgets Omni/profiler mounts, skips storage
|
||||
start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` handlers.
|
||||
"""
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_build_openai_app(args, supported_tasks):
|
||||
app = FastAPI()
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def upstream_chat():
|
||||
return {"owner": "upstream"}
|
||||
|
||||
@app.post("/v1/chat/completions/batch")
|
||||
async def upstream_batch():
|
||||
return {"owner": "upstream"}
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def upstream_models():
|
||||
return {"owner": "upstream"}
|
||||
|
||||
@app.get("/v1/upstream-only")
|
||||
async def upstream_only():
|
||||
return {"owner": "upstream"}
|
||||
|
||||
captured["supported_tasks"] = supported_tasks
|
||||
return app
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_build_async_omni(*_args, **_kwargs):
|
||||
yield _FakeEngineClient(
|
||||
stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}],
|
||||
vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)),
|
||||
)
|
||||
|
||||
async def fake_omni_init_app_state(engine_client, state, args):
|
||||
state.engine_client = engine_client
|
||||
state.initialized_by_omni = True
|
||||
|
||||
async def fake_storage_start():
|
||||
captured["storage_started"] = True
|
||||
|
||||
async def fake_serve_http(app, **_kwargs):
|
||||
captured["served_app"] = app
|
||||
task = asyncio.create_task(asyncio.sleep(0))
|
||||
await asyncio.sleep(0)
|
||||
return task
|
||||
|
||||
monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
|
||||
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
|
||||
monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
|
||||
monkeypatch.setattr(
|
||||
api_server,
|
||||
"shutdown_unsupported_routes",
|
||||
lambda app, restrictions: captured.setdefault("restrictions", restrictions),
|
||||
)
|
||||
monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start))
|
||||
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
|
||||
|
||||
sock = _FakeSocket()
|
||||
await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args())
|
||||
|
||||
served_wrapper = captured["served_app"]
|
||||
# Timestamp middleware wraps the FastAPI app.
|
||||
assert hasattr(served_wrapper, "_inner")
|
||||
served_app = served_wrapper._inner
|
||||
routes = served_app.routes
|
||||
_assert_unique_route_entries(routes)
|
||||
manifest = _route_manifest(routes)
|
||||
|
||||
assert captured["supported_tasks"] == ("generate",)
|
||||
assert captured["storage_started"] is True
|
||||
assert captured["restrictions"] == {}
|
||||
assert sock.closed is True
|
||||
assert served_app.state.initialized_by_omni is True
|
||||
|
||||
# Upstream override: unrelated route kept; Omni owns overridden pairs once.
|
||||
# Endpoint identity ensures the surviving route is the Omni handler, not
|
||||
# the fake upstream handler with the same method/path.
|
||||
assert ("GET", "/v1/upstream-only") in manifest
|
||||
assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion
|
||||
assert (
|
||||
_single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint
|
||||
is api_server.create_batch_chat_completion
|
||||
)
|
||||
assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
|
||||
|
||||
# Assembled app includes Omni router HTTP/WS surface + profiler.
|
||||
for item in _EXPECTED_ROUTER_ROUTES:
|
||||
assert item in manifest, f"assembled app missing {item}"
|
||||
assert ("POST", "/start_profile") in manifest
|
||||
assert ("POST", "/stop_profile") in manifest
|
||||
|
||||
assert served_app.exception_handlers[EngineGenerateError]
|
||||
assert served_app.exception_handlers[EngineDeadError]
|
||||
|
||||
# Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper).
|
||||
openapi_paths = served_app.openapi().get("paths", {})
|
||||
assert "/v1/chat/completions" in openapi_paths
|
||||
assert "/v1/models" in openapi_paths
|
||||
assert "/health" in openapi_paths
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None:
|
||||
"""Lock outermost timestamp middleware behavior.
|
||||
|
||||
Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes
|
||||
start getting stamped / mutated incorrectly.
|
||||
"""
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_build_openai_app(args, supported_tasks):
|
||||
return FastAPI()
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_build_async_omni(*_args, **_kwargs):
|
||||
yield _FakeEngineClient()
|
||||
|
||||
async def fake_omni_init_app_state(engine_client, state, args):
|
||||
state.engine_client = engine_client
|
||||
|
||||
async def fake_serve_http(app, **_kwargs):
|
||||
captured["served_app"] = app
|
||||
task = asyncio.create_task(asyncio.sleep(0))
|
||||
await asyncio.sleep(0)
|
||||
return task
|
||||
|
||||
monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
|
||||
monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
|
||||
monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
|
||||
monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0)))
|
||||
monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
|
||||
|
||||
await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args())
|
||||
middleware = captured["served_app"]
|
||||
|
||||
http_scope = {"type": "http", "state": {}}
|
||||
ws_scope = {"type": "websocket"}
|
||||
|
||||
async def _receive():
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
async def _send(_message):
|
||||
return None
|
||||
|
||||
called = {"http": False, "websocket": False}
|
||||
|
||||
async def inner(scope, receive, send):
|
||||
called[scope["type"]] = True
|
||||
|
||||
middleware._inner = inner
|
||||
await middleware(http_scope, _receive, _send)
|
||||
await middleware(ws_scope, _receive, _send)
|
||||
|
||||
assert called["http"] is True
|
||||
assert called["websocket"] is True
|
||||
assert "request_timestamp" in http_scope["state"]
|
||||
assert isinstance(http_scope["state"]["request_timestamp"], float)
|
||||
assert "request_timestamp" not in ws_scope.get("state", {})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "payload"),
|
||||
[
|
||||
("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}),
|
||||
("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}),
|
||||
("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}),
|
||||
("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}),
|
||||
(
|
||||
"/v1/realtime/robot/openpi",
|
||||
{"type": "error", "error": "Robot policy not available", "code": "unsupported"},
|
||||
),
|
||||
("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}),
|
||||
],
|
||||
)
|
||||
def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None:
|
||||
"""Lock missing-handler WebSocket capability probes.
|
||||
|
||||
Fails if unavailable JSON changes (keys/wording/code) or the socket stays
|
||||
open after the error frame — common when WS wrappers move between modules.
|
||||
"""
|
||||
app = FastAPI()
|
||||
app.include_router(api_server.router)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect(path) as websocket:
|
||||
assert websocket.receive_json() == payload
|
||||
with pytest.raises(WebSocketDisconnect):
|
||||
websocket.receive_text()
|
||||
|
||||
|
||||
def test_health_without_engine_returns_stable_unhealthy_response() -> None:
|
||||
"""Lock ``/health`` with no engine initialized.
|
||||
|
||||
Fails if status drifts from 503 or the unhealthy payload shape/reason
|
||||
string changes during health-route extraction.
|
||||
"""
|
||||
app = FastAPI()
|
||||
response = asyncio.run(api_server.health(_request_for(app, path="/health")))
|
||||
|
||||
assert response.status_code == 503
|
||||
assert json.loads(response.body) == {
|
||||
"status": "unhealthy",
|
||||
"reason": "No engine initialized",
|
||||
}
|
||||
|
||||
|
||||
def test_models_without_handler_returns_empty_openai_list() -> None:
|
||||
"""Lock ``/v1/models`` when ``openai_serving_models`` is unset.
|
||||
|
||||
Fails if the empty listing becomes an error, or the OpenAI list envelope
|
||||
changes (``object`` / ``data``).
|
||||
"""
|
||||
app = FastAPI()
|
||||
response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models")))
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body) == {"object": "list", "data": []}
|
||||
|
||||
|
||||
def test_images_generation_without_engine_preserves_service_unavailable_error() -> None:
|
||||
"""Lock images generation when the multi-stage engine is missing.
|
||||
|
||||
Fails if the 503 detail string changes — clients/tests use it as a
|
||||
capability signal when bootstrap failed.
|
||||
"""
|
||||
app = FastAPI()
|
||||
raw_request = _request_for(app, method="POST", path="/v1/images/generations")
|
||||
request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(api_server.generate_images(request, raw_request))
|
||||
|
||||
assert exc_info.value.status_code == 503
|
||||
assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model."
|
||||
|
||||
|
||||
def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None:
|
||||
"""Lock images generation when chat serving was not wired.
|
||||
|
||||
Fails if multi-stage images stop requiring ``openai_serving_chat``, or the
|
||||
503 detail changes after images/chat ownership splits.
|
||||
"""
|
||||
app = FastAPI()
|
||||
app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()])
|
||||
app.state.stage_configs = app.state.engine_client.stage_configs
|
||||
app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")])
|
||||
app.state.openai_serving_chat = None
|
||||
app.state.args = SimpleNamespace(max_generated_image_size=None)
|
||||
|
||||
raw_request = _request_for(app, method="POST", path="/v1/images/generations")
|
||||
request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
|
||||
monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(api_server.generate_images(request, raw_request))
|
||||
|
||||
assert exc_info.value.status_code == 503
|
||||
assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation."
|
||||
|
||||
|
||||
def test_speech_without_handler_preserves_not_found_http_error() -> None:
|
||||
"""Lock speech HTTP when ``openai_serving_speech`` is unset.
|
||||
|
||||
Fails if missing speech becomes a different status/detail after TTS route
|
||||
wrappers move out of ``api_server``.
|
||||
"""
|
||||
app = FastAPI()
|
||||
app.state.openai_serving_speech = None
|
||||
raw_request = _request_for(app, method="POST", path="/v1/audio/speech")
|
||||
request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(api_server.create_speech(request, raw_request))
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert exc_info.value.detail == "The model does not support Speech API"
|
||||
|
||||
|
||||
def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
|
||||
"""Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
|
||||
|
||||
Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers
|
||||
move (P0.2) or exception-handler wiring changes.
|
||||
|
||||
Exercises the envelope through the registered handler, not the private
|
||||
``_create_engine_error_json_response`` helper, so relocating that helper
|
||||
in P0.2 is not a false red.
|
||||
"""
|
||||
app = FastAPI()
|
||||
# TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
|
||||
# leaves api_server (red here is a rename, not a surface regression).
|
||||
api_server._register_omni_exception_handlers(app)
|
||||
app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
|
||||
app.state.server = object()
|
||||
app.state.args = SimpleNamespace(log_error_stack=False)
|
||||
|
||||
req = _request_for(app, method="POST", path="/v1/chat/completions")
|
||||
req.state.request_metadata = SimpleNamespace(request_id="req-123")
|
||||
|
||||
monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
|
||||
|
||||
exc = EngineGenerateError("boom")
|
||||
exc.error_stage_id = 2 # type: ignore[attr-defined]
|
||||
handler = app.exception_handlers[EngineGenerateError]
|
||||
response = asyncio.run(handler(req, exc))
|
||||
payload = json.loads(response.body)
|
||||
|
||||
assert "error" in payload
|
||||
assert payload["error"]["request_id"] == "req-123"
|
||||
assert payload["error"]["error_stage_id"] == 2
|
||||
assert response.status_code >= 400
|
||||
|
||||
|
||||
def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None:
|
||||
"""Lock Omni ``EngineDeadError`` handler registration + JSON envelope.
|
||||
|
||||
Fails if registration is skipped or the handler stops returning OpenAI-style
|
||||
JSON with ``request_id``.
|
||||
"""
|
||||
app = FastAPI()
|
||||
# TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
|
||||
# leaves api_server (red here is a rename, not a surface regression).
|
||||
api_server._register_omni_exception_handlers(app)
|
||||
app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
|
||||
app.state.server = object()
|
||||
app.state.args = SimpleNamespace(log_error_stack=False)
|
||||
|
||||
monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
|
||||
|
||||
handler = app.exception_handlers[EngineDeadError]
|
||||
req = _request_for(app)
|
||||
req.state.request_metadata = SimpleNamespace(request_id="dead-1")
|
||||
response = asyncio.run(handler(req, EngineDeadError()))
|
||||
payload = json.loads(response.body)
|
||||
assert payload["error"]["request_id"] == "dead-1"
|
||||
assert response.status_code >= 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
|
||||
"""Lock pure-diffusion ``app.state`` keys after init, including live vs None.
|
||||
|
||||
Fails if bootstrap drops keys route owners still read (e.g. video/speech/
|
||||
models), wires a handler that must stay None, or leaves a required
|
||||
handler as None.
|
||||
"""
|
||||
stage = SimpleNamespace(engine_args={})
|
||||
engine = _FakeEngineClient(stage_configs=[stage])
|
||||
|
||||
def _for_diffusion_factory(label: str):
|
||||
@classmethod
|
||||
def _factory(cls, *args, **kwargs):
|
||||
return _marker(label)
|
||||
|
||||
return _factory
|
||||
|
||||
monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
|
||||
monkeypatch.setattr(
|
||||
api_server.OmniOpenAIServingAudioGenerate,
|
||||
"for_diffusion",
|
||||
_for_diffusion_factory("audio_generate"),
|
||||
)
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
|
||||
monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
|
||||
monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech"))
|
||||
monkeypatch.setattr(
|
||||
api_server.ServingRealtimeRobotOpenPI,
|
||||
"create_policy_server",
|
||||
classmethod(lambda cls, *a, **k: _marker("openpi")),
|
||||
)
|
||||
|
||||
state = State()
|
||||
await api_server.omni_init_app_state(engine, state, _minimal_args())
|
||||
|
||||
_assert_app_state_snapshot(
|
||||
state,
|
||||
expected_keys=_DIFFUSION_APP_STATE_KEYS,
|
||||
must_be_wired=_DIFFUSION_MUST_BE_WIRED,
|
||||
must_be_none=_DIFFUSION_MUST_BE_NONE,
|
||||
)
|
||||
assert state.diffusion_engine is engine
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
|
||||
"""Lock multi-stage ``app.state`` keys after init, including live vs None.
|
||||
|
||||
Fails if chat/speech/video/realtime/tokenization keys disappear or are
|
||||
left as None after init refactors — classic “route still registered,
|
||||
handler never wired” failure mode.
|
||||
"""
|
||||
engine = _FakeEngineClient(
|
||||
stage_configs=[object(), object()],
|
||||
vllm_config=SimpleNamespace(
|
||||
lora_config=None,
|
||||
model_config=SimpleNamespace(),
|
||||
parallel_config=SimpleNamespace(_api_process_rank=0),
|
||||
),
|
||||
)
|
||||
|
||||
class _FakeModels:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.base_model_paths = kwargs.get("base_model_paths") or []
|
||||
|
||||
async def init_static_loras(self):
|
||||
return None
|
||||
|
||||
class _FakeCtor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def warmup(self):
|
||||
return None
|
||||
|
||||
class _FakeSpeech(_FakeCtor):
|
||||
async def warmup(self):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or [])
|
||||
monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
|
||||
monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "ServingScores", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech)
|
||||
monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video"))
|
||||
monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
|
||||
monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False)
|
||||
|
||||
state = State()
|
||||
await api_server.omni_init_app_state(engine, state, _minimal_args())
|
||||
|
||||
_assert_app_state_snapshot(
|
||||
state,
|
||||
expected_keys=_MULTISTAGE_APP_STATE_KEYS,
|
||||
must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
|
||||
must_be_none=_MULTISTAGE_MUST_BE_NONE,
|
||||
)
|
||||
@@ -4059,7 +4059,9 @@ class TestTTSAsyncOffloading:
|
||||
|
||||
assert adapter_model_type != legacy_tts_model_type
|
||||
assert any(
|
||||
call.args and call.args[0] == "TTS speech request %s: model=%s" and call.args[2] == adapter_model_type
|
||||
call.args
|
||||
and call.args[0] == "TTS speech request %s: text=%r, model=%s"
|
||||
and call.args[3] == adapter_model_type
|
||||
for call in log_info.call_args_list
|
||||
)
|
||||
|
||||
|
||||
@@ -22,10 +22,6 @@ from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
|
||||
IndexTTS25Adapter,
|
||||
indextts2_conditioning_cache_salt,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
|
||||
MossTTSAdapter,
|
||||
MossTTSNanoAdapter,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
|
||||
from vllm_omni.model_executor.models.indextts2 import prompt_utils
|
||||
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
|
||||
@@ -98,20 +94,6 @@ def test_all_adapters_are_ar_or_diffusion():
|
||||
assert cls.backend in ("ar", "diffusion")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
|
||||
def test_moss_tts_applies_request_max_new_tokens(adapter_cls):
|
||||
adapter = adapter_cls(SimpleNamespace(server=object()))
|
||||
stage_defaults = [SimpleNamespace(max_tokens=4096)]
|
||||
|
||||
overridden = adapter.apply_sampling_overrides(
|
||||
stage_defaults,
|
||||
SimpleNamespace(max_new_tokens=512),
|
||||
)
|
||||
|
||||
assert overridden[0].max_tokens == 512
|
||||
assert stage_defaults[0].max_tokens == 4096
|
||||
|
||||
|
||||
def test_qwen3_tts_metadata():
|
||||
assert Qwen3TTSAdapter.backend == "ar"
|
||||
assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
|
||||
|
||||
@@ -153,7 +153,6 @@ _PIPELINE_STAGES = [
|
||||
"indextts2_talker",
|
||||
"latent_generator",
|
||||
"llm",
|
||||
"minimax_music3_ar",
|
||||
"ming_tts",
|
||||
"moss_tts",
|
||||
"moss_tts_codec",
|
||||
|
||||
@@ -2168,31 +2168,6 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
|
||||
assert captured.quality == "high"
|
||||
|
||||
|
||||
def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture):
|
||||
_mock_encode_video_bytes(mocker)
|
||||
response = test_client.post(
|
||||
"/v1/videos/sync",
|
||||
data={
|
||||
"prompt": "drive forward",
|
||||
"extra_params": json.dumps(
|
||||
{
|
||||
"sana_wm": {"action": "d-4", "rotation_speed_deg": 1.5},
|
||||
"sana_wm_native_max_tokens": 30000,
|
||||
}
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
captured_params = engine.captured_sampling_params_list[0]
|
||||
# The camera block reaches the model through extra_args; the Sana preprocess
|
||||
# hook is what lifts it onto the prompt.
|
||||
assert captured_params.extra_args["sana_wm"]["action"] == "d-4"
|
||||
assert captured_params.extra_args["sana_wm"]["rotation_speed_deg"] == 1.5
|
||||
assert captured_params.extra_args["sana_wm_native_max_tokens"] == 30000
|
||||
|
||||
|
||||
def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
|
||||
"""Frame interpolation parameters should be forwarded on the sync path."""
|
||||
encode_mock = _mock_encode_video_bytes(mocker)
|
||||
|
||||
@@ -58,23 +58,6 @@ def test_default_stage_devices_from_sequence_parallel():
|
||||
assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
|
||||
|
||||
|
||||
def test_default_stage_devices_and_dp_from_num_gpus():
|
||||
"""Resolve omitted DP before deriving devices from an explicit GPU count."""
|
||||
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
|
||||
{
|
||||
"num_gpus": 8,
|
||||
"tensor_parallel_size": 2,
|
||||
"ulysses_degree": 2,
|
||||
}
|
||||
)[0]
|
||||
|
||||
parallel_config = stage_cfg["engine_args"]["parallel_config"]
|
||||
assert parallel_config.data_parallel_size == 2
|
||||
assert parallel_config.world_size == 8
|
||||
assert stage_cfg["engine_args"]["num_gpus"] == 8
|
||||
assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
|
||||
|
||||
|
||||
def test_default_stage_config_uses_parallel_size_kwargs():
|
||||
"""Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
|
||||
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
|
||||
@@ -93,8 +76,8 @@ def test_default_stage_config_uses_parallel_size_kwargs():
|
||||
assert parallel_config.enable_expert_parallel is True
|
||||
|
||||
|
||||
def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
|
||||
"""Keep omitted DP unresolved until the runtime WORLD size is known."""
|
||||
def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
|
||||
"""Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
|
||||
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
|
||||
{
|
||||
"pipeline_parallel_size": None,
|
||||
@@ -109,7 +92,7 @@ def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
|
||||
|
||||
parallel_config = stage_cfg["engine_args"]["parallel_config"]
|
||||
assert parallel_config.pipeline_parallel_size == 1
|
||||
assert parallel_config.data_parallel_size is None
|
||||
assert parallel_config.data_parallel_size == 1
|
||||
assert parallel_config.tensor_parallel_size == 1
|
||||
assert parallel_config.enable_expert_parallel is False
|
||||
assert stage_cfg["engine_args"]["enforce_eager"] is False
|
||||
|
||||
@@ -1,92 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for resolve_model_to_local_path."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import huggingface_hub
|
||||
import pytest
|
||||
from huggingface_hub.errors import LocalEntryNotFoundError
|
||||
|
||||
from vllm_omni.model_executor.model_loader import weight_utils
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# Both resolution paths are patched in every test, so nothing is ever fetched
|
||||
REPO_ID = "test-org/test-model"
|
||||
RESOLVED = "/cache/snapshots/deadbeef"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def paths(mocker):
|
||||
"""The resolver's two ways out, both patched."""
|
||||
return SimpleNamespace(
|
||||
cache_only=mocker.patch.object(huggingface_hub, "snapshot_download", return_value=RESOLVED),
|
||||
download=mocker.patch.object(weight_utils, "download_weights_from_hf_specific", return_value=RESOLVED),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("allow_download", [False, True])
|
||||
def test_local_directory_short_circuits(tmp_path, paths, allow_download):
|
||||
"""A local dir resolves to itself without either path running."""
|
||||
assert resolve_model_to_local_path(str(tmp_path), allow_download=allow_download) == str(tmp_path)
|
||||
|
||||
paths.cache_only.assert_not_called()
|
||||
paths.download.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "allow_patterns", "cache_dir"),
|
||||
[
|
||||
({}, None, None),
|
||||
({"allow_patterns": ["subdir/*"], "cache_dir": "/cache"}, ["subdir/*"], "/cache"),
|
||||
],
|
||||
)
|
||||
def test_repo_id_is_cache_only_by_default(paths, kwargs, allow_patterns, cache_dir):
|
||||
"""Without `allow_download` nothing is fetched, and narrowing is forwarded."""
|
||||
assert resolve_model_to_local_path(REPO_ID, **kwargs) == RESOLVED
|
||||
|
||||
paths.cache_only.assert_called_once_with(
|
||||
REPO_ID,
|
||||
cache_dir=cache_dir,
|
||||
allow_patterns=allow_patterns,
|
||||
local_files_only=True,
|
||||
)
|
||||
paths.download.assert_not_called()
|
||||
|
||||
|
||||
def test_cache_miss_propagates(paths):
|
||||
"""A failed resolution raises instead of degrading into the raw repo id."""
|
||||
paths.cache_only.side_effect = LocalEntryNotFoundError("cold cache")
|
||||
|
||||
with pytest.raises(LocalEntryNotFoundError):
|
||||
resolve_model_to_local_path(REPO_ID)
|
||||
|
||||
paths.download.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "allow_patterns"),
|
||||
[
|
||||
({}, ["*"]),
|
||||
({"allow_patterns": ["subdir/model*.safetensors"]}, ["subdir/model*.safetensors"]),
|
||||
],
|
||||
)
|
||||
def test_download_routes_through_the_locked_helper(paths, kwargs, allow_patterns):
|
||||
"""Downloads go through the vLLM helper holding the node-wide lock.
|
||||
|
||||
That helper is also what honors `VLLM_USE_MODELSCOPE`;
|
||||
calling `snapshot_download` directly would drop both.
|
||||
"""
|
||||
assert resolve_model_to_local_path(REPO_ID, allow_download=True, cache_dir="/cache", **kwargs) == RESOLVED
|
||||
|
||||
paths.download.assert_called_once_with(
|
||||
model_name_or_path=REPO_ID,
|
||||
cache_dir="/cache",
|
||||
allow_patterns=allow_patterns,
|
||||
require_all=True,
|
||||
)
|
||||
paths.cache_only.assert_not_called()
|
||||
@@ -7,14 +7,8 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
|
||||
from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
|
||||
from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler
|
||||
from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT
|
||||
from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss
|
||||
from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, get_epss_timesteps
|
||||
from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
|
||||
from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
|
||||
from vllm_omni.model_executor.models.common.ming.fm import Solver
|
||||
from vllm_omni.model_executor.models.ming_tts.constants import (
|
||||
AGGREGATOR_HIDDEN_SIZE,
|
||||
HISTORY_PATCH_SIZE,
|
||||
@@ -30,6 +24,7 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
|
||||
SPEAKER_EMBEDDING_DIM,
|
||||
VAE_PATCH_SIZE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
|
||||
from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
|
||||
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings
|
||||
from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
|
||||
@@ -49,87 +44,7 @@ def test_ming_tts_audio_vae_uses_common_config():
|
||||
|
||||
def test_ming_tts_cfm_solver_uses_common_implementation():
|
||||
"""Ming dense imports the shared solver implementation directly."""
|
||||
assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.cfm_solver"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("symbol", "expected_module"),
|
||||
[
|
||||
(Aggregator, "vllm_omni.model_executor.models.common.ming.aggregator"),
|
||||
(DiT, "vllm_omni.model_executor.models.common.ming.cfm_dit"),
|
||||
(FlowLoss, "vllm_omni.model_executor.models.common.ming.cfm_head"),
|
||||
(CFMGraphExecutor, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
|
||||
(CFMSampler, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
|
||||
(SpeakerEmbeddingExtractor, "vllm_omni.model_executor.models.common.ming.speaker_extractor"),
|
||||
(SpkembExtractor, "vllm_omni.model_executor.models.common.ming.spk_embedding"),
|
||||
],
|
||||
)
|
||||
def test_ming_tts_reusable_modules_live_in_common(symbol, expected_module):
|
||||
"""The pieces a second Ming-family TTS model reuses live under common/ming.
|
||||
|
||||
Guards the promotion: if an implementation drifts back into `ming_tts`
|
||||
(or into another model package), a second model would have to reach across
|
||||
packages to get it.
|
||||
"""
|
||||
assert symbol.__module__ == expected_module
|
||||
|
||||
|
||||
def test_ming_flash_talker_uses_the_shared_aggregator():
|
||||
"""Ming flash no longer carries its own copy of the aggregator.
|
||||
|
||||
Both Ming packages ran near-identical aggregators one directory apart,
|
||||
which is the duplication this module layout exists to remove.
|
||||
"""
|
||||
from vllm_omni.model_executor.models.ming_flash_omni import talker_module
|
||||
|
||||
assert talker_module.Aggregator is Aggregator
|
||||
assert not any(
|
||||
isinstance(obj, type) and obj.__module__ == talker_module.__name__ and name == "Aggregator"
|
||||
for name, obj in vars(talker_module).items()
|
||||
)
|
||||
|
||||
|
||||
def test_shared_aggregator_matches_the_ming_flash_construction():
|
||||
"""The shared aggregator reproduces Ming flash's former default geometry.
|
||||
|
||||
`talker_module.Aggregator` defaulted to `in_channels=64` while the
|
||||
shared class defaults to 4, so the call site now passes `latent_dim`
|
||||
explicitly. This pins the resulting shapes rather than the default.
|
||||
"""
|
||||
latent_dim, llm_hidden = 64, 32
|
||||
aggregator = Aggregator(
|
||||
in_channels=latent_dim,
|
||||
llm_input_dim=llm_hidden,
|
||||
hidden_size=16,
|
||||
depth=1,
|
||||
num_heads=2,
|
||||
)
|
||||
|
||||
assert aggregator.x_embedder.in_features == latent_dim
|
||||
assert aggregator.final_layer.linear.out_features == llm_hidden
|
||||
|
||||
out = aggregator(torch.randn(1, 3, latent_dim))
|
||||
# Keeps only the CLS projection: [B, T, D] -> [B, 1, llm_hidden].
|
||||
assert out.shape == (1, 1, llm_hidden)
|
||||
|
||||
|
||||
def test_shared_aggregator_rejects_mismatched_latent_width():
|
||||
"""Wrong-width input fails loudly instead of silently projecting incorrect."""
|
||||
aggregator = Aggregator(in_channels=64, llm_input_dim=32, hidden_size=16, depth=1, num_heads=2)
|
||||
|
||||
with pytest.raises(ValueError, match="feature dim mismatch"):
|
||||
aggregator(torch.randn(1, 3, 4))
|
||||
|
||||
|
||||
def test_cfm_solver_owns_the_timestep_schedule():
|
||||
"""get_epss_timesteps sits with the solver, not with the DiT blocks."""
|
||||
from vllm_omni.model_executor.models.common.ming import dit_blocks
|
||||
|
||||
assert not hasattr(dit_blocks, "get_epss_timesteps")
|
||||
|
||||
t = get_epss_timesteps(10, device="cpu", dtype=torch.float32)
|
||||
assert t.shape == (11,)
|
||||
assert float(t[0]) == 0.0 and float(t[-1]) == 1.0
|
||||
assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
|
||||
|
||||
|
||||
def test_ming_tts_flowloss_preserves_checkpoint_prefix():
|
||||
|
||||
@@ -1,156 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression tests for MiniCPM-o 4.5 streaming audio cache compatibility."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.models.whisper.modeling_whisper import WhisperConfig
|
||||
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
|
||||
WHISPER_ATTENTION_CLASSES,
|
||||
MiniCPMWhisperEncoder,
|
||||
MiniCPMWhisperEncoderLayer,
|
||||
_get_audio_cache_length,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _LegacyWhisperAttention(nn.Module):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__()
|
||||
self.seen_cache = None
|
||||
|
||||
def forward(self, hidden_states, past_key_value=None, **kwargs):
|
||||
self.seen_cache = past_key_value
|
||||
return hidden_states, None, past_key_value
|
||||
|
||||
|
||||
class _CurrentWhisperAttention(nn.Module):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__()
|
||||
self.seen_cache = None
|
||||
|
||||
def forward(self, hidden_states, past_key_values=None, **kwargs):
|
||||
self.seen_cache = past_key_values
|
||||
return hidden_states, None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("implementation", "attention_cls"),
|
||||
[
|
||||
("test_legacy_cache", _LegacyWhisperAttention),
|
||||
("test_current_cache", _CurrentWhisperAttention),
|
||||
],
|
||||
)
|
||||
def test_whisper_attention_preserves_streaming_cache(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
implementation: str,
|
||||
attention_cls: type[nn.Module],
|
||||
) -> None:
|
||||
monkeypatch.setitem(WHISPER_ATTENTION_CLASSES, implementation, attention_cls)
|
||||
config = WhisperConfig(
|
||||
d_model=4,
|
||||
encoder_attention_heads=1,
|
||||
encoder_ffn_dim=8,
|
||||
encoder_layers=1,
|
||||
)
|
||||
config._attn_implementation = implementation
|
||||
layer = MiniCPMWhisperEncoderLayer(config, layer_idx=0)
|
||||
cache = object()
|
||||
|
||||
outputs = layer(
|
||||
torch.randn(1, 2, config.d_model),
|
||||
attention_mask=None,
|
||||
layer_head_mask=None,
|
||||
past_key_values=cache,
|
||||
use_cache=True,
|
||||
)
|
||||
|
||||
assert layer.self_attn.seen_cache is cache
|
||||
assert outputs[-1] is cache
|
||||
|
||||
|
||||
def test_audio_cache_length_uses_current_cache_api() -> None:
|
||||
dynamic_cache = SimpleNamespace(get_seq_length=lambda: 7)
|
||||
encoder_decoder_cache = SimpleNamespace(self_attention_cache=dynamic_cache)
|
||||
|
||||
assert _get_audio_cache_length(encoder_decoder_cache) == 7
|
||||
|
||||
|
||||
def test_audio_cache_length_supports_legacy_cache() -> None:
|
||||
key = torch.zeros(1, 1, 5, 2)
|
||||
value = torch.zeros_like(key)
|
||||
|
||||
assert _get_audio_cache_length(((key, value),)) == 5
|
||||
|
||||
|
||||
def _tiny_encoder_config(attn_implementation: str) -> WhisperConfig:
|
||||
config = WhisperConfig(
|
||||
num_mel_bins=16,
|
||||
d_model=32,
|
||||
encoder_layers=2,
|
||||
encoder_attention_heads=2,
|
||||
encoder_ffn_dim=64,
|
||||
max_source_positions=100,
|
||||
dropout=0.0,
|
||||
attention_dropout=0.0,
|
||||
)
|
||||
config._attn_implementation = attn_implementation
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"])
|
||||
def test_encoder_reuses_cache_across_streaming_chunks(attn_implementation: str) -> None:
|
||||
"""Encode two chunks against the installed Whisper attention and cache classes.
|
||||
|
||||
Unlike the stub-based tests above, this drives the real transformers cache
|
||||
API that ``get_audio_embedding_streaming`` relies on, so it fails whenever a
|
||||
removed cache method or renamed attention kwarg breaks streaming reuse.
|
||||
"""
|
||||
torch.manual_seed(0)
|
||||
config = _tiny_encoder_config(attn_implementation)
|
||||
encoder = MiniCPMWhisperEncoder(config).eval()
|
||||
|
||||
first_frames, second_frames = 20, 14
|
||||
first_len = (first_frames - 1) // 2 + 1
|
||||
second_len = (second_frames - 1) // 2 + 1
|
||||
first_mel = torch.randn(1, config.num_mel_bins, first_frames)
|
||||
second_mel = torch.randn(1, config.num_mel_bins, second_frames)
|
||||
|
||||
with torch.no_grad():
|
||||
first = encoder(
|
||||
first_mel,
|
||||
attention_mask=torch.zeros(1, 1, first_len, first_len),
|
||||
past_key_values=None,
|
||||
use_cache=True,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
cache = first.past_key_values
|
||||
assert cache is not None
|
||||
assert _get_audio_cache_length(cache) == first_len
|
||||
|
||||
second = encoder(
|
||||
second_mel,
|
||||
attention_mask=torch.zeros(1, 1, second_len, first_len + second_len),
|
||||
past_key_values=cache,
|
||||
use_cache=True,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
assert _get_audio_cache_length(second.past_key_values) == first_len + second_len
|
||||
|
||||
stateless = encoder(
|
||||
second_mel,
|
||||
attention_mask=torch.zeros(1, 1, second_len, second_len),
|
||||
past_key_values=None,
|
||||
use_cache=True,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
|
||||
assert torch.isfinite(second.last_hidden_state).all()
|
||||
assert not torch.allclose(second.last_hidden_state, stateless.last_hidden_state)
|
||||
@@ -1,41 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""CPU regression tests for MOSS-TTS compact top-k sampling."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
|
||||
|
||||
|
||||
def test_compact_topk_only_samples_retained_candidates() -> None:
|
||||
logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(7))
|
||||
|
||||
sampled = modeling_moss_tts_local._sample_token(
|
||||
logits,
|
||||
temperature=1.7,
|
||||
top_k=25,
|
||||
top_p=0.8,
|
||||
do_sample=True,
|
||||
generator=torch.Generator().manual_seed(42),
|
||||
)
|
||||
|
||||
retained = logits.topk(25, dim=-1).indices
|
||||
assert torch.all((retained == sampled.unsqueeze(-1)).any(dim=-1))
|
||||
|
||||
|
||||
def test_compact_topk_tiny_nucleus_selects_argmax() -> None:
|
||||
logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(17))
|
||||
|
||||
sampled = modeling_moss_tts_local._sample_token(
|
||||
logits,
|
||||
temperature=1.7,
|
||||
top_k=25,
|
||||
top_p=1e-9,
|
||||
do_sample=True,
|
||||
generator=torch.Generator().manual_seed(42),
|
||||
)
|
||||
|
||||
torch.testing.assert_close(sampled, logits.argmax(dim=-1))
|
||||
@@ -1,44 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import GPT2Config
|
||||
|
||||
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
|
||||
MossTTSLocalDepthTransformer,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
|
||||
|
||||
|
||||
def _test_config() -> GPT2Config:
|
||||
return GPT2Config(
|
||||
n_embd=32,
|
||||
n_head=4,
|
||||
n_inner=64,
|
||||
layer_norm_epsilon=1e-6,
|
||||
rope_base=10000.0,
|
||||
)
|
||||
|
||||
|
||||
def test_rope_cache_matches_reference_and_reuses_storage() -> None:
|
||||
model = MossTTSLocalDepthTransformer(_test_config()).eval()
|
||||
attn = model.h[0].attn
|
||||
device = torch.device("cpu")
|
||||
|
||||
attn.prepare_rope_cache(12, device, torch.float32)
|
||||
cache_ptr = attn._rope_cos_cache.data_ptr()
|
||||
cos, sin = attn._rope_cos_sin(4, device, torch.float32)
|
||||
|
||||
positions = torch.arange(4, dtype=torch.float32)
|
||||
freqs = torch.einsum("s,d->sd", positions, attn.inv_freq.float())
|
||||
expected_cos = freqs.cos().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
|
||||
expected_sin = freqs.sin().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
|
||||
torch.testing.assert_close(cos, expected_cos)
|
||||
torch.testing.assert_close(sin, expected_sin)
|
||||
|
||||
# Smaller slices retain the fixed allocation used by CUDA Graph.
|
||||
attn.prepare_rope_cache(6, device, torch.float32)
|
||||
assert attn._rope_cos_cache.data_ptr() == cache_ptr
|
||||
@@ -23,11 +23,6 @@ class _RecordingLogitsProcessor:
|
||||
return torch.zeros((1, 4))
|
||||
|
||||
|
||||
class _EmbeddingBackbone(nn.Module):
|
||||
def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return input_ids.to(dtype=torch.float32).unsqueeze(-1).repeat(1, 4)
|
||||
|
||||
|
||||
def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None:
|
||||
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
|
||||
MossTTSDelayTalkerForGeneration,
|
||||
@@ -46,74 +41,3 @@ def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> N
|
||||
assert logits.shape == (1, 4)
|
||||
assert model.logits_processor.args == (model.text_lm_head, hidden_states)
|
||||
assert model.logits_processor.kwargs == {}
|
||||
|
||||
|
||||
def test_moss_tts_local_keeps_fixed_shape_rows_and_routes_stop_logits() -> None:
|
||||
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
|
||||
MossTTSLocalTalkerForGeneration,
|
||||
)
|
||||
|
||||
model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
|
||||
nn.Module.__init__(model)
|
||||
model.audio_pad_token_id = 1024
|
||||
model.audio_assistant_slot_token_id = 2
|
||||
model.im_end_token_id = 3
|
||||
model.text_vocab_size = 8
|
||||
model.n_vq = 4
|
||||
model._batch_state = None
|
||||
model._batch_state_spans = None
|
||||
model._batch_should_continue = None
|
||||
|
||||
valid = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)
|
||||
stopped = torch.full((1, 4), model.audio_pad_token_id, dtype=torch.long)
|
||||
infos = [
|
||||
{"audio_state": {}, "audio_codes": {"current": valid}},
|
||||
{"audio_state": {}, "audio_codes": {"current": stopped}},
|
||||
]
|
||||
spans = [(0, 3), (3, 4)]
|
||||
|
||||
output = model.make_omni_output(
|
||||
torch.randn(4, 8),
|
||||
model_intermediate_buffer=infos,
|
||||
request_token_spans=spans,
|
||||
)
|
||||
|
||||
audio = output.multimodal_outputs["codes"]["audio"]
|
||||
assert len(audio) == 2
|
||||
torch.testing.assert_close(audio[0], valid)
|
||||
# Keeping this row avoids CUDA boolean-index shape discovery; the
|
||||
# talker2codec raw async-chunk processor drops all-pad rows on CPU.
|
||||
torch.testing.assert_close(audio[1], stopped)
|
||||
assert model._batch_should_continue.tolist() == [True, False]
|
||||
|
||||
logits = model.compute_logits(torch.randn(4, 8))
|
||||
|
||||
assert logits is not None
|
||||
assert logits.argmax(dim=-1).tolist() == [2, 2, 2, 3]
|
||||
|
||||
|
||||
def test_moss_tts_local_single_token_prefill_does_not_advance_audio() -> None:
|
||||
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
|
||||
MossTTSLocalTalkerForGeneration,
|
||||
)
|
||||
|
||||
model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
|
||||
nn.Module.__init__(model)
|
||||
model.model = _EmbeddingBackbone()
|
||||
model.n_vq = 4
|
||||
|
||||
input_ids = torch.tensor([7], dtype=torch.long)
|
||||
_, embeds, update = model.preprocess(
|
||||
input_ids,
|
||||
None,
|
||||
_omni_is_prefill=True,
|
||||
audio_state={"is_stopping": True},
|
||||
ref_offset=5,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(embeds, torch.full((1, 4), 7.0))
|
||||
assert update == {
|
||||
"audio_state": {"is_stopping": False},
|
||||
"ref_offset": 6,
|
||||
}
|
||||
assert "mtp_inputs" not in update
|
||||
@@ -1,132 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression tests for MiniMax Music 3 batched payload partitioning.
|
||||
|
||||
The stage-0 payload is split per request by indexing list values: a list is
|
||||
indexed by batch row, anything else is handed to every request unchanged. So
|
||||
per-request metadata has to be carried as lists. Merging it into one dict made
|
||||
every request in a batch receive the last request's seed, window offset and
|
||||
terminal flag, which crosses seeds between songs, misplaces acoustic windows,
|
||||
and lets one request's terminal flag end another request's stream early.
|
||||
|
||||
These drive ``make_omni_output`` directly with a hand-built batch, so no GPU
|
||||
and no engine are needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.minimax_music3.talker import (
|
||||
MiniMaxMusic3TalkerForConditionalGeneration as Talker,
|
||||
)
|
||||
from vllm_omni.utils.mm_outputs import to_payload_element
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_WIDTH = 4
|
||||
|
||||
|
||||
def _info(*, seed: int, start: int, frames: int, last: bool) -> dict:
|
||||
"""One batch row's intermediate buffer, as ``postprocess`` leaves it."""
|
||||
return {
|
||||
"latent": torch.full((frames, _WIDTH), float(seed)),
|
||||
"codes": {"audio": torch.full((frames, 8), seed)},
|
||||
"meta": {
|
||||
"audio_seed": seed,
|
||||
"num_processed_tokens": start,
|
||||
"codec_chunk_frames": frames,
|
||||
"last_chunk": last,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _payload_for(output, row: int) -> dict:
|
||||
"""Partition the batch payload down to one request, as the runner does."""
|
||||
flat = {}
|
||||
for key, value in output.multimodal_outputs.items():
|
||||
if isinstance(value, dict):
|
||||
for qual, val in value.items():
|
||||
flat[f"{key}.{qual}"] = val
|
||||
else:
|
||||
flat[key] = value
|
||||
return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()}
|
||||
|
||||
|
||||
class TestBatchedMetadataPartitioning:
|
||||
"""Two requests in one batch must keep their own metadata."""
|
||||
|
||||
@staticmethod
|
||||
def _two_request_batch():
|
||||
infos = [
|
||||
_info(seed=111, start=100, frames=3, last=False),
|
||||
_info(seed=222, start=200, frames=5, last=True),
|
||||
]
|
||||
talker = Talker.__new__(Talker)
|
||||
return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
|
||||
|
||||
def test_seeds_are_not_crossed(self) -> None:
|
||||
"""A song decoded with another request's seed is a different song."""
|
||||
output = self._two_request_batch()
|
||||
assert _payload_for(output, 0)["meta.audio_seed"] == 111
|
||||
assert _payload_for(output, 1)["meta.audio_seed"] == 222
|
||||
|
||||
def test_window_offsets_stay_with_their_request(self) -> None:
|
||||
"""The start frame decides where a window lands in the song."""
|
||||
output = self._two_request_batch()
|
||||
assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100
|
||||
assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200
|
||||
|
||||
def test_terminal_flag_does_not_leak_across_requests(self) -> None:
|
||||
"""One request finishing must not end another request's stream."""
|
||||
output = self._two_request_batch()
|
||||
assert _payload_for(output, 0)["meta.last_chunk"] is False
|
||||
assert _payload_for(output, 1)["meta.last_chunk"] is True
|
||||
|
||||
def test_span_length_matches_its_own_latent(self) -> None:
|
||||
output = self._two_request_batch()
|
||||
for row, frames in ((0, 3), (1, 5)):
|
||||
entry = _payload_for(output, row)
|
||||
assert entry["meta.codec_chunk_frames"] == frames
|
||||
assert entry["latent"].shape[0] == frames
|
||||
|
||||
def test_latent_and_codes_follow_the_same_row(self) -> None:
|
||||
"""The tensors were already per-request; pin that they stay aligned."""
|
||||
output = self._two_request_batch()
|
||||
for row, seed in ((0, 111), (1, 222)):
|
||||
entry = _payload_for(output, row)
|
||||
assert entry["latent"][0, 0].item() == pytest.approx(float(seed))
|
||||
assert int(entry["codes.audio"][0, 0]) == seed
|
||||
|
||||
|
||||
class TestPartialBatches:
|
||||
"""Only some rows carry a span on any given step."""
|
||||
|
||||
def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None:
|
||||
infos = [{}, _info(seed=777, start=300, frames=2, last=True)]
|
||||
talker = Talker.__new__(Talker)
|
||||
output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
|
||||
|
||||
idle = _payload_for(output, 0)
|
||||
assert idle["latent"].numel() == 0
|
||||
assert idle["meta.audio_seed"] is None
|
||||
assert idle["meta.last_chunk"] is None
|
||||
|
||||
active = _payload_for(output, 1)
|
||||
assert active["meta.audio_seed"] == 777
|
||||
assert active["meta.last_chunk"] is True
|
||||
|
||||
def test_batch_with_no_spans_emits_no_payload(self) -> None:
|
||||
talker = Talker.__new__(Talker)
|
||||
output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}])
|
||||
assert output.multimodal_outputs is None
|
||||
|
||||
def test_consumed_entries_are_cleared_from_the_buffer(self) -> None:
|
||||
"""The runner never clears this buffer, so a span left behind would be
|
||||
re-emitted on every later step."""
|
||||
infos = [_info(seed=1, start=0, frames=2, last=False)]
|
||||
talker = Talker.__new__(Talker)
|
||||
talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos)
|
||||
assert "latent" not in infos[0]
|
||||
assert "meta" not in infos[0]
|
||||
@@ -1,214 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping.
|
||||
|
||||
The talker emits overlapping windows while dropping frames it can no longer
|
||||
need, and flushes whatever is left as one span when the song ends. Getting the
|
||||
eviction boundary wrong by one frame silently corrupts the conditioning handed
|
||||
to the acoustic stage from the second window onward, which is audible but hard
|
||||
to attribute, so the lifecycle is simulated here frame by frame.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows
|
||||
from vllm_omni.model_executor.models.minimax_music3.constants import (
|
||||
AR_CHUNK_FRAMES,
|
||||
AR_CHUNK_HOP_FRAMES,
|
||||
NUM_CODEBOOKS,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_music3.frame_state import (
|
||||
FrameBuffer,
|
||||
RequestARState,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# One scalar per frame instead of 32768 dims: the bookkeeping is what is under
|
||||
# test, and a frame's value doubles as its identity.
|
||||
_WIDTH = 4
|
||||
|
||||
|
||||
def _frame(index: int) -> torch.Tensor:
|
||||
return torch.full((1, _WIDTH), float(index))
|
||||
|
||||
|
||||
def _identities(chunk: torch.Tensor) -> list[int]:
|
||||
return [int(row[0].item()) for row in chunk]
|
||||
|
||||
|
||||
class TestFrameBuffer:
|
||||
def test_slice_returns_the_requested_frames(self) -> None:
|
||||
buffer = FrameBuffer()
|
||||
for index in range(10):
|
||||
buffer.append(_frame(index))
|
||||
assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6]
|
||||
|
||||
def test_slice_is_relative_to_the_live_base(self) -> None:
|
||||
buffer = FrameBuffer()
|
||||
for index in range(10):
|
||||
buffer.append(_frame(index))
|
||||
buffer.discard_before(4)
|
||||
assert buffer.base_frame == 4
|
||||
assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7]
|
||||
|
||||
def test_slicing_a_discarded_range_raises(self) -> None:
|
||||
buffer = FrameBuffer()
|
||||
for index in range(10):
|
||||
buffer.append(_frame(index))
|
||||
buffer.discard_before(5)
|
||||
with pytest.raises(RuntimeError, match="unavailable"):
|
||||
buffer.slice(3, 6)
|
||||
|
||||
def test_slicing_past_the_end_raises(self) -> None:
|
||||
buffer = FrameBuffer()
|
||||
for index in range(4):
|
||||
buffer.append(_frame(index))
|
||||
with pytest.raises(RuntimeError, match="unavailable"):
|
||||
buffer.slice(2, 9)
|
||||
|
||||
def test_discard_is_clamped_at_both_ends(self) -> None:
|
||||
buffer = FrameBuffer()
|
||||
for index in range(5):
|
||||
buffer.append(_frame(index))
|
||||
buffer.discard_before(-3)
|
||||
assert buffer.base_frame == 0
|
||||
buffer.discard_before(99)
|
||||
assert buffer.base_frame == 5
|
||||
assert buffer.end_frame == 5
|
||||
|
||||
def test_peak_tracks_the_high_water_mark(self) -> None:
|
||||
buffer = FrameBuffer()
|
||||
for index in range(10):
|
||||
buffer.append(_frame(index))
|
||||
buffer.discard_before(8)
|
||||
assert buffer.peak_frames == 10
|
||||
|
||||
|
||||
class TestWindowLifecycle:
|
||||
"""Drive the state exactly as the talker does, one frame per step."""
|
||||
|
||||
@staticmethod
|
||||
def _run(total_frames: int) -> tuple[list[list[int]], int]:
|
||||
"""Generate ``total_frames`` frames, returning the emitted spans.
|
||||
|
||||
Returns:
|
||||
The frame identities of every emitted span, and the peak buffer
|
||||
occupancy reached along the way.
|
||||
"""
|
||||
state = RequestARState(seed=1, max_frames=total_frames)
|
||||
spans: list[list[int]] = []
|
||||
for index in range(total_frames):
|
||||
state.frames.append(_frame(index))
|
||||
state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
|
||||
state.generated_frames += 1
|
||||
window = state.ready_window()
|
||||
if window is not None:
|
||||
spans.append(_identities(state.take_window(window, is_final=False)))
|
||||
|
||||
# Terminal flush: everything not yet emitted, as one contiguous span.
|
||||
start = state.emitted_windows * AR_CHUNK_HOP_FRAMES
|
||||
end = state.frames.end_frame
|
||||
if end > start:
|
||||
spans.append(_identities(state.frames.slice(start, end)))
|
||||
return spans, state.frames.peak_frames
|
||||
|
||||
def test_short_song_is_one_span(self) -> None:
|
||||
spans, _ = self._run(50)
|
||||
assert spans == [list(range(50))]
|
||||
|
||||
def test_every_frame_is_delivered_exactly_once_in_order(self) -> None:
|
||||
"""Spans overlap by design, but their union must be the whole song.
|
||||
|
||||
Reconstruct by taking each span's non-overlapping remainder, which is
|
||||
what the acoustic stage does after cropping.
|
||||
"""
|
||||
for total in (250, 437, 600, 901):
|
||||
spans, _ = self._run(total)
|
||||
rebuilt: list[int] = []
|
||||
for span in spans:
|
||||
rebuilt.extend(frame for frame in span if frame >= len(rebuilt))
|
||||
assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames"
|
||||
|
||||
def test_windows_advance_by_the_hop(self) -> None:
|
||||
spans, _ = self._run(600)
|
||||
for index, span in enumerate(spans[:-1]):
|
||||
assert span[0] == index * AR_CHUNK_HOP_FRAMES
|
||||
assert len(span) == AR_CHUNK_FRAMES
|
||||
|
||||
def test_mid_stream_windows_are_never_the_final_frame(self) -> None:
|
||||
"""A window is held back until a further frame proves it is not last."""
|
||||
total = 600
|
||||
spans, _ = self._run(total)
|
||||
for span in spans[:-1]:
|
||||
assert span[-1] < total - 1
|
||||
|
||||
def test_terminal_span_covers_the_tail(self) -> None:
|
||||
total = 437
|
||||
spans, _ = self._run(total)
|
||||
assert spans[-1][-1] == total - 1
|
||||
|
||||
def test_terminal_span_is_bounded(self) -> None:
|
||||
"""The tail is at most one full window plus one hop."""
|
||||
for total in (250, 301, 437, 600, 901, 1000):
|
||||
spans, _ = self._run(total)
|
||||
assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
|
||||
|
||||
def test_buffer_stays_bounded_over_a_long_song(self) -> None:
|
||||
"""Peak occupancy must not grow with song length; that is the point
|
||||
of evicting retired frames."""
|
||||
_, peak_short = self._run(600)
|
||||
_, peak_long = self._run(3000)
|
||||
assert peak_long == peak_short
|
||||
assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
|
||||
|
||||
def test_emitted_window_count_matches_the_window_plan(self) -> None:
|
||||
"""What the stream emitted must reconcile with chunk_windows()."""
|
||||
for total in (250, 437, 600):
|
||||
spans, _ = self._run(total)
|
||||
planned = chunk_windows(total)
|
||||
# Every window is accounted for: the pre-terminal spans are one
|
||||
# window each, and the terminal span carries whatever remains.
|
||||
assert len(spans) <= len(planned)
|
||||
assert spans[0][0] == planned[0].start
|
||||
|
||||
|
||||
class TestSamplingPositions:
|
||||
def test_each_frame_reserves_one_slot_per_codebook(self) -> None:
|
||||
"""Codebook streams must never collide, or codes repeat across frames."""
|
||||
state = RequestARState(seed=1, max_frames=10)
|
||||
positions = [state.advance_position() for _ in range(5)]
|
||||
assert positions == [0, 8, 16, 24, 32]
|
||||
|
||||
def test_draw_slots_are_globally_unique(self) -> None:
|
||||
state = RequestARState(seed=1, max_frames=64)
|
||||
slots: set[int] = set()
|
||||
for _ in range(64):
|
||||
base = state.advance_position()
|
||||
for codebook in range(NUM_CODEBOOKS):
|
||||
slot = base + codebook
|
||||
assert slot not in slots
|
||||
slots.add(slot)
|
||||
assert len(slots) == 64 * NUM_CODEBOOKS
|
||||
|
||||
|
||||
class TestCodeRetention:
|
||||
def test_codes_span_matches_the_frame_span(self) -> None:
|
||||
state = RequestARState(seed=1, max_frames=10)
|
||||
for index in range(10):
|
||||
state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
|
||||
codes = state.take_codes(2, 6)
|
||||
assert codes is not None
|
||||
assert codes.shape == (4, NUM_CODEBOOKS)
|
||||
assert [int(row[0]) for row in codes] == [2, 3, 4, 5]
|
||||
|
||||
def test_out_of_range_span_is_dropped_not_raised(self) -> None:
|
||||
"""Codes are diagnostic; the audio path must not fail without them."""
|
||||
state = RequestARState(seed=1, max_frames=10)
|
||||
for index in range(4):
|
||||
state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
|
||||
assert state.take_codes(2, 99) is None
|
||||
assert state.take_codes(0, 0) is None
|
||||
@@ -1,203 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic.
|
||||
|
||||
The prompt transformations are not cosmetic: they decide token positions, and
|
||||
token positions decide the audio. These tests pin the documented behaviours
|
||||
that a well-meaning cleanup would otherwise "fix".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.minimax_music3.chunking import (
|
||||
chunk_windows,
|
||||
crop_sample_bounds,
|
||||
overlap_mel_length,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_music3.constants import (
|
||||
AR_CHUNK_FRAMES,
|
||||
AR_CHUNK_HOP_FRAMES,
|
||||
BLEND_MEL_FRAMES,
|
||||
DAV_HOP_SAMPLES,
|
||||
HOP_MEL_FRAMES,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_music3.prompt import (
|
||||
SPECIAL_TOKEN_IDS,
|
||||
build_cfg_null_token_ids,
|
||||
build_prompt,
|
||||
clean_caption,
|
||||
normalize_lyrics,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestNormalizeLyrics:
|
||||
def test_prepends_start_tag(self) -> None:
|
||||
assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n")
|
||||
|
||||
def test_lowercases_structure_tags(self) -> None:
|
||||
assert "[verse]" in normalize_lyrics("[Verse]\nhello")
|
||||
assert "[Verse]" not in normalize_lyrics("[Verse]\nhello")
|
||||
|
||||
def test_text_after_a_leading_tag_is_dropped(self) -> None:
|
||||
"""Documented contract: a tag must sit on its own line.
|
||||
|
||||
Lyrics written next to a tag are silently lost, which is why the
|
||||
serving docs tell callers to put tags on their own line.
|
||||
"""
|
||||
out = normalize_lyrics("[Verse] Walking down the street")
|
||||
assert "Walking" not in out
|
||||
assert "[verse]" in out
|
||||
|
||||
def test_text_on_its_own_line_survives(self) -> None:
|
||||
out = normalize_lyrics("[Verse]\nWalking down the street")
|
||||
assert "Walking down the street" in out
|
||||
|
||||
def test_bracket_space_becomes_newline(self) -> None:
|
||||
assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]"
|
||||
|
||||
|
||||
class TestCleanCaption:
|
||||
def test_special_token_form_is_rewritten_not_passed_through(self) -> None:
|
||||
"""A caption must not be able to inject a real special token."""
|
||||
assert clean_caption("<|bpm 92|>") == "bpm is 92"
|
||||
assert "<|" not in clean_caption("<|bpm 92|>")
|
||||
|
||||
def test_lone_tag_keeps_its_inner_text(self) -> None:
|
||||
assert clean_caption("<|instrumental|>") == "instrumental"
|
||||
|
||||
def test_markdown_is_stripped(self) -> None:
|
||||
assert clean_caption("## Heading") == "Heading"
|
||||
assert clean_caption("**bold**") == "bold"
|
||||
assert clean_caption("- item") == "item"
|
||||
|
||||
def test_blank_lines_collapse_at_two_not_three(self) -> None:
|
||||
"""The source collapses runs of two or more, not three or more.
|
||||
|
||||
This shifts blank-line token positions and therefore the audio, so it
|
||||
is pinned deliberately.
|
||||
"""
|
||||
assert clean_caption("a\n\nb") == "a\nb"
|
||||
assert clean_caption("a\n\n\n\nb") == "a\nb"
|
||||
|
||||
|
||||
class TestBuildPrompt:
|
||||
def test_framing_tokens_are_present_and_ordered(self) -> None:
|
||||
prompt = build_prompt("a bright pop song", "[Verse]\nhello")
|
||||
assert prompt.startswith("<|im_start|><|caption_start|>")
|
||||
assert prompt.endswith("<|im_end|><|audio_start|>")
|
||||
assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>")
|
||||
|
||||
|
||||
class TestBuildCfgNullTokenIds:
|
||||
@staticmethod
|
||||
def _prompt_ids(body_len: int) -> list[int]:
|
||||
return (
|
||||
[SPECIAL_TOKEN_IDS["<|im_start|>"]]
|
||||
+ list(range(1000, 1000 + body_len))
|
||||
+ [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]]
|
||||
)
|
||||
|
||||
def test_length_is_preserved(self) -> None:
|
||||
"""Equal length is what lets the two rows decode in lockstep."""
|
||||
ids = self._prompt_ids(12)
|
||||
assert len(build_cfg_null_token_ids(ids)) == len(ids)
|
||||
|
||||
def test_body_is_replaced_and_framing_survives(self) -> None:
|
||||
ids = self._prompt_ids(12)
|
||||
null = build_cfg_null_token_ids(ids)
|
||||
assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"]
|
||||
assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"]
|
||||
assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"]
|
||||
assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]}
|
||||
|
||||
def test_input_is_not_mutated(self) -> None:
|
||||
ids = self._prompt_ids(6)
|
||||
before = list(ids)
|
||||
build_cfg_null_token_ids(ids)
|
||||
assert ids == before
|
||||
|
||||
def test_minimal_prompt_with_empty_body(self) -> None:
|
||||
ids = [
|
||||
SPECIAL_TOKEN_IDS["<|im_start|>"],
|
||||
SPECIAL_TOKEN_IDS["<|audio_cfg|>"],
|
||||
SPECIAL_TOKEN_IDS["<|im_end|>"],
|
||||
SPECIAL_TOKEN_IDS["<|audio_start|>"],
|
||||
]
|
||||
assert build_cfg_null_token_ids(ids) == ids
|
||||
|
||||
def test_too_short_is_rejected(self) -> None:
|
||||
with pytest.raises(ValueError, match="too short"):
|
||||
build_cfg_null_token_ids([1, 2, 3])
|
||||
|
||||
def test_wrong_framing_is_rejected(self) -> None:
|
||||
with pytest.raises(ValueError, match="framing"):
|
||||
build_cfg_null_token_ids([9, 9, 9, 9, 9])
|
||||
|
||||
|
||||
class TestChunkWindows:
|
||||
def test_no_frames_yields_nothing(self) -> None:
|
||||
assert chunk_windows(0) == []
|
||||
|
||||
def test_short_song_is_one_final_window(self) -> None:
|
||||
windows = chunk_windows(50)
|
||||
assert len(windows) == 1
|
||||
assert (windows[0].start, windows[0].end) == (0, 50)
|
||||
assert windows[0].is_first and windows[0].is_last
|
||||
|
||||
def test_exactly_one_window(self) -> None:
|
||||
windows = chunk_windows(AR_CHUNK_FRAMES)
|
||||
assert len(windows) == 1
|
||||
assert windows[0].is_last
|
||||
|
||||
def test_windows_advance_by_the_hop(self) -> None:
|
||||
windows = chunk_windows(450)
|
||||
starts = [w.start for w in windows]
|
||||
assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES))
|
||||
assert windows[-1].is_last
|
||||
assert not any(w.is_last for w in windows[:-1])
|
||||
|
||||
def test_windows_cover_every_frame(self) -> None:
|
||||
frames = 437
|
||||
windows = chunk_windows(frames)
|
||||
covered: set[int] = set()
|
||||
for window in windows:
|
||||
covered.update(range(window.start, window.end))
|
||||
assert covered == set(range(frames))
|
||||
|
||||
def test_final_window_is_clamped_to_the_song(self) -> None:
|
||||
windows = chunk_windows(437)
|
||||
assert windows[-1].end == 437
|
||||
|
||||
def test_only_the_first_window_is_first(self) -> None:
|
||||
windows = chunk_windows(600)
|
||||
assert windows[0].is_first
|
||||
assert not any(w.is_first for w in windows[1:])
|
||||
|
||||
|
||||
class TestCropBounds:
|
||||
def test_single_window_song_is_not_cropped(self) -> None:
|
||||
window = chunk_windows(50)[0]
|
||||
assert crop_sample_bounds(window) == (0, 0)
|
||||
|
||||
def test_first_window_keeps_its_head(self) -> None:
|
||||
left, right = crop_sample_bounds(chunk_windows(600)[0])
|
||||
assert left == 0
|
||||
assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES
|
||||
|
||||
def test_last_window_keeps_its_tail(self) -> None:
|
||||
left, right = crop_sample_bounds(chunk_windows(600)[-1])
|
||||
assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
|
||||
assert right == 0
|
||||
|
||||
def test_middle_window_is_cropped_on_both_sides(self) -> None:
|
||||
windows = chunk_windows(900)
|
||||
middle = windows[len(windows) // 2]
|
||||
left, right = crop_sample_bounds(middle)
|
||||
assert left > 0 and right > 0
|
||||
|
||||
def test_overlap_is_half_a_window(self) -> None:
|
||||
assert overlap_mel_length() == HOP_MEL_FRAMES // 2
|
||||
@@ -1,217 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression tests for MiniMax Music 3 row-to-request binding.
|
||||
|
||||
vLLM compacts and reorders the persistent batch whenever a request finishes or
|
||||
is preempted, so a batch row number is only meaningful within the step that
|
||||
produced it. Every failure pinned here was a real defect: state keyed by row
|
||||
survived into a step where the row belonged to someone else, or state keyed by
|
||||
one spelling of the request id was released under another and leaked.
|
||||
|
||||
These exercise the pure helpers and the binding arithmetic directly, so they
|
||||
need no GPU and no engine.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.minimax_music3.constants import (
|
||||
DEFAULT_MAX_AUDIO_FRAMES,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_music3.talker import (
|
||||
_CFG_UNCOND_SUFFIX,
|
||||
_first_int,
|
||||
_pair_id_from_request,
|
||||
_request_key,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestRequestKey:
|
||||
def test_engine_list_form_is_unwrapped(self) -> None:
|
||||
"""The engine injects global_request_id as a single-element list.
|
||||
|
||||
Keying state on the list's ``str()`` would produce "['req-1']", which
|
||||
never matches the bare id that ``on_requests_finished`` delivers, so
|
||||
every finished song would leak its frame buffer.
|
||||
"""
|
||||
assert _request_key(["req-1"]) == "req-1"
|
||||
|
||||
def test_bare_string_is_unchanged(self) -> None:
|
||||
assert _request_key("req-1") == "req-1"
|
||||
|
||||
def test_list_and_bare_forms_agree(self) -> None:
|
||||
"""Allocation and release must land on the same key."""
|
||||
assert _request_key(["req-1"]) == _request_key("req-1")
|
||||
|
||||
def test_empty_and_missing_are_stable(self) -> None:
|
||||
assert _request_key(None) == ""
|
||||
assert _request_key([]) == ""
|
||||
|
||||
def test_tuple_form(self) -> None:
|
||||
assert _request_key(("req-9",)) == "req-9"
|
||||
|
||||
|
||||
class TestPairIdDerivation:
|
||||
def test_conditioned_request_is_its_own_pair(self) -> None:
|
||||
assert _pair_id_from_request("req-1") == "req-1"
|
||||
|
||||
def test_companion_maps_to_its_partner(self) -> None:
|
||||
assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1"
|
||||
|
||||
def test_pair_members_agree_on_the_pair_id(self) -> None:
|
||||
cond = "req-42"
|
||||
uncond = f"req-42{_CFG_UNCOND_SUFFIX}"
|
||||
assert _pair_id_from_request(cond) == _pair_id_from_request(uncond)
|
||||
|
||||
def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None:
|
||||
assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond"
|
||||
|
||||
|
||||
class TestFirstInt:
|
||||
def test_serving_layer_list_wrapping(self) -> None:
|
||||
"""tts_params values arrive wrapped in a single-element list."""
|
||||
assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750
|
||||
|
||||
def test_bare_value(self) -> None:
|
||||
assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750
|
||||
|
||||
def test_missing_falls_back(self) -> None:
|
||||
assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
|
||||
assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
|
||||
|
||||
def test_nonsense_falls_back_rather_than_raising(self) -> None:
|
||||
assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
|
||||
assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
|
||||
assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
|
||||
|
||||
|
||||
def _feedback_plan(
|
||||
request_ids: list[str],
|
||||
computed: list[int],
|
||||
scheduled: list[int],
|
||||
prompt_tokens: dict[str, int],
|
||||
has_codes: set[str],
|
||||
) -> list[int]:
|
||||
"""Mirror of the talker's feedback row selection.
|
||||
|
||||
Returns the flat token offsets whose embedding gets replaced by an audio
|
||||
frame. Kept as a standalone function so the arithmetic can be checked
|
||||
without constructing a model.
|
||||
"""
|
||||
offsets: list[int] = []
|
||||
running = 0
|
||||
for count in scheduled:
|
||||
offsets.append(running)
|
||||
running += count
|
||||
|
||||
rows: list[int] = []
|
||||
for row, request_id in enumerate(request_ids):
|
||||
if request_id not in has_codes:
|
||||
continue
|
||||
if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0):
|
||||
continue
|
||||
rows.append(offsets[row])
|
||||
return rows
|
||||
|
||||
|
||||
class TestFeedbackRowSelection:
|
||||
"""The guard that decides which token embeddings become audio frames."""
|
||||
|
||||
def test_decode_only_step_replaces_every_row(self) -> None:
|
||||
ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"]
|
||||
rows = _feedback_plan(
|
||||
ids,
|
||||
computed=[120, 120, 90, 90],
|
||||
scheduled=[1, 1, 1, 1],
|
||||
prompt_tokens={i: 50 for i in ids},
|
||||
has_codes=set(ids),
|
||||
)
|
||||
assert rows == [0, 1, 2, 3]
|
||||
|
||||
def test_prefill_row_is_never_overwritten(self) -> None:
|
||||
"""A prompt token must keep its real embedding.
|
||||
|
||||
The earlier guard compared the flat token count against a row budget,
|
||||
so a short prompt slipped through and had its first embeddings
|
||||
replaced with a previous request's audio frame.
|
||||
"""
|
||||
rows = _feedback_plan(
|
||||
["a", "b"],
|
||||
computed=[0, 300],
|
||||
scheduled=[8, 1],
|
||||
prompt_tokens={"a": 8, "b": 40},
|
||||
has_codes={"b"},
|
||||
)
|
||||
# Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single
|
||||
# decode token sits at flat offset 8.
|
||||
assert rows == [8]
|
||||
|
||||
def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None:
|
||||
"""The exact shape that defeated the old token-count guard."""
|
||||
rows = _feedback_plan(
|
||||
["a"],
|
||||
computed=[0],
|
||||
scheduled=[10],
|
||||
prompt_tokens={"a": 10},
|
||||
has_codes={"a"},
|
||||
)
|
||||
assert rows == []
|
||||
|
||||
def test_chunked_prefill_midway_is_not_fed_back(self) -> None:
|
||||
rows = _feedback_plan(
|
||||
["a"],
|
||||
computed=[2048],
|
||||
scheduled=[2048],
|
||||
prompt_tokens={"a": 5000},
|
||||
has_codes={"a"},
|
||||
)
|
||||
assert rows == []
|
||||
|
||||
def test_final_prefill_chunk_is_not_fed_back(self) -> None:
|
||||
"""The step that completes a prompt still carries prompt tokens."""
|
||||
rows = _feedback_plan(
|
||||
["a"],
|
||||
computed=[4096],
|
||||
scheduled=[904],
|
||||
prompt_tokens={"a": 5000},
|
||||
has_codes={"a"},
|
||||
)
|
||||
assert rows == []
|
||||
|
||||
def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None:
|
||||
"""Tokens are a flat concatenation, so offsets are a running sum."""
|
||||
ids = ["p", "d1", "d2"]
|
||||
rows = _feedback_plan(
|
||||
ids,
|
||||
computed=[0, 500, 700],
|
||||
scheduled=[64, 1, 1],
|
||||
prompt_tokens={"p": 64, "d1": 40, "d2": 40},
|
||||
has_codes={"d1", "d2"},
|
||||
)
|
||||
assert rows == [64, 65]
|
||||
|
||||
def test_request_without_codes_yet_is_skipped(self) -> None:
|
||||
"""The first decode step after prefill has no previous frame."""
|
||||
rows = _feedback_plan(
|
||||
["a"],
|
||||
computed=[50],
|
||||
scheduled=[1],
|
||||
prompt_tokens={"a": 50},
|
||||
has_codes=set(),
|
||||
)
|
||||
assert rows == []
|
||||
|
||||
def test_reordered_batch_binds_by_request_not_position(self) -> None:
|
||||
"""The same two requests in swapped rows get swapped offsets.
|
||||
|
||||
This is the compaction case: holding codes in a row-indexed buffer
|
||||
would have handed each request the other's frame.
|
||||
"""
|
||||
prompt_tokens = {"a": 50, "b": 50}
|
||||
first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"})
|
||||
swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"})
|
||||
assert first == [0]
|
||||
assert swapped == [1]
|
||||
@@ -1,389 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for the model-neutral CFG request-pairing scheduler patches.
|
||||
|
||||
The patches keep a cond/uncond request pair step-locked inside one engine:
|
||||
admitted together, kept at equal ``num_computed_tokens``, finished together.
|
||||
They are installed onto a lightweight fake scheduler here; vLLM's real
|
||||
``Scheduler`` needs a full engine (and a GPU) to construct.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.common import cfg_pairing
|
||||
from vllm_omni.model_executor.models.common.cfg_pairing import (
|
||||
_equalize_cfg_pair_progress,
|
||||
_patch_scheduler,
|
||||
cfg_uncond_suffixes,
|
||||
register_cfg_uncond_suffix,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
# ───────── fakes ─────────
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
"""Minimal stand-in for ``vllm.v1.request.Request`` (identity equality)."""
|
||||
|
||||
def __init__(self, request_id: str, extra_args: dict[str, Any] | None = None, step_tokens: int = 1) -> None:
|
||||
self.request_id = request_id
|
||||
self.sampling_params = SimpleNamespace(extra_args=extra_args)
|
||||
self.num_computed_tokens = 0
|
||||
self.step_tokens = step_tokens
|
||||
self.skip_reading_prefix_cache = False
|
||||
|
||||
|
||||
class _FakeWaitingQueue(deque):
|
||||
"""FCFS waiting queue: vLLM's is also a ``deque`` subclass."""
|
||||
|
||||
def remove_requests(self, requests: list[_FakeRequest]) -> None:
|
||||
keep = [req for req in self if not any(req is dropped for dropped in requests)]
|
||||
self.clear()
|
||||
self.extend(keep)
|
||||
|
||||
def prepend_request(self, request: _FakeRequest) -> None:
|
||||
self.appendleft(request)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeSchedulerOutput:
|
||||
num_scheduled_tokens: dict[str, int] = field(default_factory=dict)
|
||||
total_num_scheduled_tokens: int = 0
|
||||
|
||||
|
||||
def _make_scheduler_cls() -> type:
|
||||
"""Fresh class per test so the in-place patches never stack."""
|
||||
|
||||
class FakeScheduler:
|
||||
def __init__(self) -> None:
|
||||
self.waiting = _FakeWaitingQueue()
|
||||
self.running: list[_FakeRequest] = []
|
||||
self.requests: dict[str, _FakeRequest] = {}
|
||||
self.finished: list[set[str]] = []
|
||||
self.observed_waiting: list[list[str]] = []
|
||||
self.max_num_scheduled_tokens = 512
|
||||
self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=4096)
|
||||
|
||||
def add_request(self, request: _FakeRequest) -> None:
|
||||
self.requests[request.request_id] = request
|
||||
self.waiting.append(request)
|
||||
|
||||
def finish_requests(self, request_ids: Any, finished_status: Any) -> None:
|
||||
if request_ids is None:
|
||||
ids = set(self.requests)
|
||||
elif isinstance(request_ids, str):
|
||||
ids = {request_ids}
|
||||
else:
|
||||
ids = set(request_ids)
|
||||
self.finished.append(ids)
|
||||
self.running = [req for req in self.running if req.request_id not in ids]
|
||||
for req_id in ids:
|
||||
self.requests.pop(req_id, None)
|
||||
|
||||
def schedule(self) -> _FakeSchedulerOutput:
|
||||
self.observed_waiting.append([req.request_id for req in self.waiting])
|
||||
while self.waiting:
|
||||
self.running.append(self.waiting.popleft())
|
||||
output = _FakeSchedulerOutput()
|
||||
for req in self.running:
|
||||
req.num_computed_tokens += req.step_tokens
|
||||
output.num_scheduled_tokens[req.request_id] = req.step_tokens
|
||||
output.total_num_scheduled_tokens += req.step_tokens
|
||||
return output
|
||||
|
||||
return FakeScheduler
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scheduler() -> Any:
|
||||
cls = _make_scheduler_cls()
|
||||
_patch_scheduler(cls)
|
||||
return cls()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _restore_suffix_registry():
|
||||
original = set(cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES)
|
||||
yield
|
||||
cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.clear()
|
||||
cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.update(original)
|
||||
|
||||
|
||||
def _add_pair(
|
||||
scheduler: Any,
|
||||
pair_id: str | None = None,
|
||||
base_id: str = "req-1",
|
||||
suffix: str = "__cfg_uncond",
|
||||
) -> tuple[_FakeRequest, _FakeRequest]:
|
||||
cond_extra: dict[str, Any] = {"cfg_role": "cond"}
|
||||
uncond_extra: dict[str, Any] = {"cfg_role": "uncond"}
|
||||
if pair_id is not None:
|
||||
cond_extra["cfg_pair_id"] = pair_id
|
||||
uncond_extra["cfg_pair_id"] = pair_id
|
||||
cond = _FakeRequest(base_id, cond_extra)
|
||||
uncond = _FakeRequest(f"{base_id}{suffix}", uncond_extra)
|
||||
scheduler.add_request(cond)
|
||||
scheduler.add_request(uncond)
|
||||
return cond, uncond
|
||||
|
||||
|
||||
# ───────── pair-id derivation ─────────
|
||||
|
||||
|
||||
def test_pair_id_is_derived_from_request_id_when_absent(scheduler: Any) -> None:
|
||||
cond, uncond = _add_pair(scheduler)
|
||||
|
||||
assert scheduler._cfg_pairs == {"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}}
|
||||
assert scheduler._cfg_req_to_pair == {cond.request_id: "req-1", uncond.request_id: "req-1"}
|
||||
|
||||
|
||||
def test_registered_companion_suffix_is_stripped(scheduler: Any) -> None:
|
||||
register_cfg_uncond_suffix("::negative")
|
||||
assert "::negative" in cfg_uncond_suffixes()
|
||||
|
||||
cond, uncond = _add_pair(scheduler, base_id="music-7", suffix="::negative")
|
||||
|
||||
assert scheduler._cfg_pairs == {"music-7": {"cond": cond.request_id, "uncond": uncond.request_id}}
|
||||
|
||||
|
||||
def test_register_rejects_empty_suffix() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
register_cfg_uncond_suffix("")
|
||||
|
||||
|
||||
def test_uncond_without_known_suffix_keeps_its_own_request_id(scheduler: Any) -> None:
|
||||
# Degenerate input: the companion never joins the cond pair, so the pair
|
||||
# stays incomplete and the hold budget eventually releases it.
|
||||
scheduler.add_request(_FakeRequest("solo-uncond", {"cfg_role": "uncond"}))
|
||||
|
||||
assert scheduler._cfg_pairs == {"solo-uncond": {"uncond": "solo-uncond"}}
|
||||
|
||||
|
||||
def test_explicit_pair_id_takes_precedence(scheduler: Any) -> None:
|
||||
cond, uncond = _add_pair(scheduler, pair_id="pair-A")
|
||||
|
||||
assert scheduler._cfg_pairs == {"pair-A": {"cond": cond.request_id, "uncond": uncond.request_id}}
|
||||
assert cond.request_id not in scheduler._cfg_pairs
|
||||
assert scheduler._cfg_req_to_pair[uncond.request_id] == "pair-A"
|
||||
|
||||
|
||||
def test_cfg_requests_skip_prefix_cache_reads(scheduler: Any) -> None:
|
||||
cond, uncond = _add_pair(scheduler)
|
||||
|
||||
assert cond.skip_reading_prefix_cache is True
|
||||
assert uncond.skip_reading_prefix_cache is True
|
||||
|
||||
|
||||
# ───────── admission holding ─────────
|
||||
|
||||
|
||||
def test_lone_member_is_held_then_released(scheduler: Any) -> None:
|
||||
scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
|
||||
|
||||
scheduler.schedule()
|
||||
|
||||
assert scheduler.observed_waiting == [[]] # the base scheduler saw nothing
|
||||
assert [req.request_id for req in scheduler.waiting] == ["req-1"]
|
||||
assert scheduler.running == []
|
||||
|
||||
|
||||
def test_complete_pair_is_admitted_together(scheduler: Any) -> None:
|
||||
scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
|
||||
scheduler.schedule()
|
||||
scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
|
||||
|
||||
scheduler.schedule()
|
||||
|
||||
assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond"]
|
||||
assert [req.request_id for req in scheduler.running] == ["req-1", "req-1__cfg_uncond"]
|
||||
assert list(scheduler.waiting) == []
|
||||
|
||||
|
||||
def test_hold_budget_expiry_drops_the_stale_pair(scheduler: Any, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(cfg_pairing, "_MAX_PAIR_HOLD_STEPS", 2)
|
||||
scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
|
||||
|
||||
scheduler.schedule()
|
||||
scheduler.schedule()
|
||||
assert scheduler.running == []
|
||||
assert scheduler._cfg_pairs # still hoping for the partner
|
||||
|
||||
scheduler.schedule()
|
||||
|
||||
assert scheduler._cfg_pairs == {}
|
||||
assert scheduler._cfg_req_to_pair == {}
|
||||
assert [req.request_id for req in scheduler.running] == ["req-1"]
|
||||
|
||||
|
||||
def test_waiting_queue_reorder_puts_partners_adjacent(scheduler: Any) -> None:
|
||||
scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
|
||||
scheduler.add_request(_FakeRequest("other", None))
|
||||
scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
|
||||
|
||||
scheduler.schedule()
|
||||
|
||||
assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond", "other"]
|
||||
|
||||
|
||||
# ───────── progress equalization ─────────
|
||||
|
||||
|
||||
def _equalize_fixture(
|
||||
cond_computed: int,
|
||||
uncond_computed: int,
|
||||
cond_sched: int,
|
||||
uncond_sched: int,
|
||||
) -> tuple[Any, _FakeSchedulerOutput, _FakeRequest, _FakeRequest]:
|
||||
cond = _FakeRequest("req-1", {"cfg_role": "cond"})
|
||||
uncond = _FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})
|
||||
cond.num_computed_tokens = cond_computed
|
||||
uncond.num_computed_tokens = uncond_computed
|
||||
fake = SimpleNamespace(
|
||||
_cfg_pairs={"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}},
|
||||
requests={cond.request_id: cond, uncond.request_id: uncond},
|
||||
)
|
||||
output = _FakeSchedulerOutput(
|
||||
num_scheduled_tokens={cond.request_id: cond_sched, uncond.request_id: uncond_sched},
|
||||
total_num_scheduled_tokens=cond_sched + uncond_sched,
|
||||
)
|
||||
return fake, output, cond, uncond
|
||||
|
||||
|
||||
def test_equalize_pulls_the_ahead_member_back() -> None:
|
||||
fake, output, cond, uncond = _equalize_fixture(10, 6, 5, 4)
|
||||
|
||||
_equalize_cfg_pair_progress(fake, output)
|
||||
|
||||
assert cond.num_computed_tokens == 6
|
||||
assert uncond.num_computed_tokens == 6
|
||||
assert output.num_scheduled_tokens == {cond.request_id: 1, uncond.request_id: 4}
|
||||
assert output.total_num_scheduled_tokens == 5
|
||||
|
||||
|
||||
def test_equalize_is_a_noop_when_the_pair_is_already_aligned() -> None:
|
||||
fake, output, cond, uncond = _equalize_fixture(8, 8, 4, 4)
|
||||
|
||||
_equalize_cfg_pair_progress(fake, output)
|
||||
|
||||
assert (cond.num_computed_tokens, uncond.num_computed_tokens) == (8, 8)
|
||||
assert output.total_num_scheduled_tokens == 8
|
||||
|
||||
|
||||
def test_equalize_skips_an_infeasible_rollback() -> None:
|
||||
# Clawing back 4 tokens would leave the cond row with 0 scheduled tokens,
|
||||
# which the scheduler output cannot express; leave the pair to _drop_split_pairs.
|
||||
fake, output, cond, uncond = _equalize_fixture(10, 6, 4, 4)
|
||||
|
||||
_equalize_cfg_pair_progress(fake, output)
|
||||
|
||||
assert cond.num_computed_tokens == 10
|
||||
assert output.total_num_scheduled_tokens == 8
|
||||
|
||||
|
||||
# ───────── finish ─────────
|
||||
|
||||
|
||||
def test_finish_requests_finishes_the_partner(scheduler: Any) -> None:
|
||||
cond, uncond = _add_pair(scheduler)
|
||||
scheduler.schedule()
|
||||
|
||||
scheduler.finish_requests(cond.request_id, "FINISHED_STOPPED")
|
||||
|
||||
assert scheduler.finished == [{cond.request_id, uncond.request_id}]
|
||||
assert scheduler._cfg_pairs == {}
|
||||
assert scheduler._cfg_req_to_pair == {}
|
||||
assert scheduler.running == []
|
||||
|
||||
|
||||
def test_finish_all_requests_clears_the_registry(scheduler: Any) -> None:
|
||||
_add_pair(scheduler)
|
||||
scheduler.schedule()
|
||||
|
||||
scheduler.finish_requests(None, "FINISHED_ABORTED")
|
||||
|
||||
assert scheduler._cfg_pairs == {}
|
||||
assert scheduler._cfg_req_to_pair == {}
|
||||
|
||||
|
||||
# ───────── non-CFG traffic ─────────
|
||||
|
||||
|
||||
def test_plain_requests_are_untouched(scheduler: Any) -> None:
|
||||
scheduler.add_request(_FakeRequest("plain-1", None))
|
||||
scheduler.add_request(_FakeRequest("plain-2", {"seed": 3}))
|
||||
|
||||
output = scheduler.schedule()
|
||||
|
||||
assert scheduler._cfg_pairs == {}
|
||||
assert scheduler.observed_waiting == [["plain-1", "plain-2"]]
|
||||
assert output.num_scheduled_tokens == {"plain-1": 1, "plain-2": 1}
|
||||
assert output.total_num_scheduled_tokens == 2
|
||||
# No pair ever existed, so the hold bookkeeping is never allocated.
|
||||
assert not hasattr(scheduler, "_cfg_hold_counts")
|
||||
|
||||
|
||||
def test_default_suffix_matches_the_audex_stage_processor() -> None:
|
||||
from vllm_omni.model_executor.stage_input_processors.audex import CFG_UNCOND_SUFFIX
|
||||
|
||||
assert CFG_UNCOND_SUFFIX in cfg_uncond_suffixes()
|
||||
|
||||
|
||||
def test_audex_cfg_module_still_exports_the_patch_entry_points() -> None:
|
||||
from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches, cfg_patches_applied
|
||||
|
||||
assert apply_cfg_patches is cfg_pairing.apply_cfg_patches
|
||||
assert cfg_patches_applied is cfg_pairing.cfg_patches_applied
|
||||
|
||||
|
||||
# ───────── install gate ─────────
|
||||
|
||||
|
||||
def test_gate_triggers_on_a_cfg_logits_processor() -> None:
|
||||
from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
|
||||
|
||||
class AudexCFGLogitsProcessor:
|
||||
pass
|
||||
|
||||
assert _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[AudexCFGLogitsProcessor]))
|
||||
# vLLM keeps unresolved entries as dotted strings.
|
||||
assert _stage_declares_cfg_pairs(
|
||||
SimpleNamespace(logits_processors=["vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"])
|
||||
)
|
||||
|
||||
|
||||
def test_gate_triggers_on_a_declared_cfg_role_default() -> None:
|
||||
from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
|
||||
|
||||
assert _stage_declares_cfg_pairs(
|
||||
SimpleNamespace(logits_processors=None, sampling_extra_args_keys=("cfg_role", "cfg_scale"))
|
||||
)
|
||||
|
||||
|
||||
def test_gate_stays_off_for_non_cfg_stages() -> None:
|
||||
from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
|
||||
|
||||
assert not _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[], sampling_extra_args_keys=()))
|
||||
assert not _stage_declares_cfg_pairs(
|
||||
SimpleNamespace(logits_processors=["some.module:MinPLogitsProcessor"], sampling_extra_args_keys=("tta_rvq",))
|
||||
)
|
||||
assert not _stage_declares_cfg_pairs(SimpleNamespace())
|
||||
|
||||
|
||||
def test_sampling_extra_args_keys_are_derived_from_stage_defaults() -> None:
|
||||
from vllm_omni.engine.stage_init_utils import _sampling_extra_args_keys
|
||||
|
||||
assert _sampling_extra_args_keys({"extra_args": {"cfg_scale": 1.5, "cfg_role": "cond"}}) == (
|
||||
"cfg_role",
|
||||
"cfg_scale",
|
||||
)
|
||||
assert _sampling_extra_args_keys({"temperature": 0.1}) == ()
|
||||
assert _sampling_extra_args_keys(None) == ()
|
||||
@@ -80,7 +80,6 @@ EXCLUDED_MODELS = [
|
||||
"Pi0Pipeline",
|
||||
"SoulXSingerPipeline",
|
||||
"SoulXSingerSVCPipeline",
|
||||
"SanaWmPipeline",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@
|
||||
#
|
||||
# stability (when included in TEST_TYPE):
|
||||
# From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py
|
||||
# model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22 + hunyuan_image; all → all five
|
||||
# model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six
|
||||
# LABEL_SUBSTR: if set, script path / job key / filename must contain it
|
||||
#
|
||||
# Requirements: bash, python3, PyYAML (pip install pyyaml)
|
||||
@@ -588,6 +588,7 @@ def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -
|
||||
STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [
|
||||
("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)),
|
||||
("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)),
|
||||
("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)),
|
||||
("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)),
|
||||
("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)),
|
||||
("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)),
|
||||
|
||||
@@ -7,8 +7,7 @@ Dataset source: https://huggingface.co/datasets/liarliar/Daily-Omni
|
||||
|
||||
Supports loading QA metadata from:
|
||||
- Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
|
||||
- HuggingFace Hub id (``dataset_path``): resolved with ``snapshot_download`` to the repo's
|
||||
``qa.json``, which reads the hub cache and therefore also works under ``HF_HUB_OFFLINE``
|
||||
- HuggingFace datasets (``dataset_path``): legacy online mode
|
||||
|
||||
Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir``
|
||||
is not set, the first request that needs on-disk media downloads that archive from the Hugging Face
|
||||
@@ -25,9 +24,9 @@ Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``?
|
||||
This class therefore inherits only ``BenchmarkDataset`` (minimal: ``dataset_path``,
|
||||
``random_seed``, ``self.data``) and implements **two explicit loaders**:
|
||||
``_load_from_local_json`` (default path for air-gapped runs) and ``_load_from_huggingface``
|
||||
(Hub ids, resolved to the repo's ``qa.json`` through the hub cache). The latter is **not**
|
||||
(optional legacy path for users who prefer ``datasets`` + Hub cache). The latter is **not**
|
||||
inheritance; it is the same Hub rows as before, factored into a helper so one class can
|
||||
serve both deployment modes without mandatory ``datasets``.
|
||||
serve both deployment modes without mandatory ``datasets`` when using ``qa_json_path``.
|
||||
|
||||
Usage:
|
||||
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset
|
||||
@@ -39,7 +38,7 @@ Usage:
|
||||
random_seed=42,
|
||||
)
|
||||
|
||||
# Hub id mode (works offline once the repo is in the HF hub cache)
|
||||
# HuggingFace mode (legacy, requires network)
|
||||
dataset = DailyOmniDataset(
|
||||
dataset_path="liarliar/Daily-Omni",
|
||||
dataset_split="train",
|
||||
@@ -211,68 +210,6 @@ def _resolve_hf_cache_snapshot(path: Path) -> Path:
|
||||
return path
|
||||
|
||||
|
||||
#: QA metadata pulled by :func:`ensure_daily_omni_hub_root`. ``Videos.tar`` is deliberately left
|
||||
#: out: it is ~60GB and only :func:`ensure_daily_omni_hub_videos_dir` needs it, lazily, once a
|
||||
#: request actually asks for media.
|
||||
_DAILY_OMNI_QA_ALLOW_PATTERNS = ["qa.json", "QA.json"]
|
||||
|
||||
|
||||
def ensure_daily_omni_hub_root(repo_id: str) -> Path:
|
||||
"""Resolve a Daily-Omni Hub id to its snapshot directory, downloading QA metadata if needed.
|
||||
|
||||
Mirrors :func:`...seed_tts_dataset.resolve_seed_tts_root`: ``snapshot_download`` resolves the
|
||||
repo through the **hub cache**, so an air-gapped run (``HF_HUB_OFFLINE=1``) keeps working as
|
||||
long as the dataset was pre-staged there. ``datasets.load_dataset`` cannot do this — it only
|
||||
consults the separate *datasets* cache, which a ``hf download`` never populates.
|
||||
|
||||
Raises:
|
||||
ImportError: if ``huggingface_hub`` is not installed.
|
||||
ValueError: if ``repo_id`` is empty.
|
||||
"""
|
||||
rid = (repo_id or "").strip()
|
||||
if not rid:
|
||||
raise ValueError("repo_id is required to download Daily-Omni from Hugging Face")
|
||||
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Install huggingface_hub to load Daily-Omni from the Hub, or pass a local "
|
||||
"--daily-omni-qa-json / --daily-omni-video-dir mirror."
|
||||
) from e
|
||||
|
||||
cache = snapshot_download(
|
||||
repo_id=rid,
|
||||
repo_type="dataset",
|
||||
allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS,
|
||||
)
|
||||
root = _resolve_hf_cache_snapshot(Path(cache).resolve())
|
||||
logger.info("Daily-Omni Hub snapshot ready at %s (repo=%s)", root, rid)
|
||||
return root
|
||||
|
||||
|
||||
def _daily_omni_qa_load_error(
|
||||
repo_id: str | None,
|
||||
snapshot_error: Exception | None,
|
||||
datasets_error: Exception,
|
||||
) -> RuntimeError:
|
||||
"""Compose an actionable error after both QA-metadata sources failed.
|
||||
|
||||
Offline CI hits this when the dataset is missing from the hub cache *or* staged under a
|
||||
different repo id, so the message names both attempts and the flags that bypass them.
|
||||
"""
|
||||
parts = [f"Could not load Daily-Omni QA metadata for {repo_id!r}."]
|
||||
if snapshot_error is not None:
|
||||
parts.append(f"huggingface_hub snapshot_download failed: {type(snapshot_error).__name__}: {snapshot_error}")
|
||||
parts.append(f"datasets.load_dataset failed: {type(datasets_error).__name__}: {datasets_error}")
|
||||
parts.append(
|
||||
"For offline runs, pre-stage the dataset in the HF hub cache "
|
||||
"(`hf download --repo-type dataset <repo>`) and make sure the repo id matches, "
|
||||
"or pass --daily-omni-qa-json / --daily-omni-video-dir pointing at a local mirror."
|
||||
)
|
||||
return RuntimeError(" ".join(parts))
|
||||
|
||||
|
||||
def resolve_daily_omni_local_root(dataset_path: str | None) -> Path | None:
|
||||
"""Return the local Daily-Omni mirror root for ``dataset_path``, or ``None`` if not local.
|
||||
|
||||
@@ -496,8 +433,7 @@ class DailyOmniDataset(BenchmarkDataset):
|
||||
|
||||
QA metadata can be loaded from:
|
||||
- Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
|
||||
- HuggingFace Hub id (``dataset_path``): the repo's ``qa.json`` via ``snapshot_download``
|
||||
(hub-cache backed, offline-safe); ``datasets.load_dataset`` only as a parquet fallback
|
||||
- HuggingFace datasets (``dataset_path``): legacy online mode
|
||||
|
||||
Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set,
|
||||
the first sample that needs on-disk media downloads that archive from the Hugging Face dataset
|
||||
@@ -508,8 +444,7 @@ class DailyOmniDataset(BenchmarkDataset):
|
||||
qa_json_path: Path to local qa.json file (offline mode, preferred). When provided,
|
||||
``dataset_path`` and ``dataset_split`` are ignored.
|
||||
dataset_path: HuggingFace dataset path (e.g., "liarliar/Daily-Omni"). Used only if
|
||||
``qa_json_path`` is not provided; resolved through the hub cache, so a pre-staged
|
||||
repo also loads with ``HF_HUB_OFFLINE=1``.
|
||||
``qa_json_path`` is not provided (legacy online mode).
|
||||
dataset_split: Dataset split to use (default: "train"). Used only in online mode.
|
||||
random_seed: Random seed for shuffling
|
||||
video_dir: Directory containing extracted video files (default: None; may be filled lazily
|
||||
@@ -667,42 +602,21 @@ class DailyOmniDataset(BenchmarkDataset):
|
||||
self.data = _ListDatasetIterator(data)
|
||||
|
||||
def _load_from_huggingface(self) -> None:
|
||||
"""Load QA rows for a Hub id, preferring the repo's ``qa.json`` over ``datasets``.
|
||||
"""Load QA rows via ``datasets.load_dataset`` (legacy / convenience path).
|
||||
|
||||
The Hub repo ships ``qa.json`` (1,197 rows) next to ``Videos.tar``, so the primary path is
|
||||
:func:`ensure_daily_omni_hub_root` → :meth:`_load_from_local_json`: ``snapshot_download``
|
||||
reads the hub cache and therefore works under ``HF_HUB_OFFLINE``, the same contract as
|
||||
Seed-TTS. ``datasets.load_dataset`` remains only as a fallback for mirrors
|
||||
that publish parquet rows instead of ``qa.json``; it consults the *datasets* cache rather
|
||||
than the hub cache and streams shards over HTTP, so it cannot serve air-gapped runs.
|
||||
Kept for backward compatibility: callers can still pass ``dataset_path=liarliar/Daily-Omni``
|
||||
and get the same parquet-backed rows as the Hub dataset card, with streaming (or
|
||||
non-streaming if ``no_stream=True``) and shuffle.
|
||||
|
||||
This is intentionally **not** implemented by subclassing ``HuggingFaceDataset``: that base
|
||||
always runs Hub ``load_dataset`` from its constructor and expects a Hub id as the primary
|
||||
API; Daily-Omni instead chooses the source in ``load_data()`` (JSON vs Hub) while sharing
|
||||
one ``sample()`` / request-building implementation for both.
|
||||
"""
|
||||
snapshot_error: Exception | None = None
|
||||
try:
|
||||
root = ensure_daily_omni_hub_root(self.dataset_path)
|
||||
except Exception as e: # noqa: BLE001 - any failure just falls through to `datasets`
|
||||
snapshot_error = e
|
||||
else:
|
||||
qa_json = daily_omni_local_qa_json(root)
|
||||
if qa_json is not None:
|
||||
self.qa_json_path = qa_json
|
||||
logger.info("Loading Daily-Omni QA metadata from Hub snapshot: %s", qa_json)
|
||||
self._load_from_local_json()
|
||||
return
|
||||
logger.info(
|
||||
"Hub snapshot %s ships no qa.json; falling back to `datasets.load_dataset`",
|
||||
root,
|
||||
)
|
||||
|
||||
if load_dataset is None:
|
||||
raise _daily_omni_qa_load_error(
|
||||
self.dataset_path,
|
||||
snapshot_error,
|
||||
ImportError("datasets library is not installed (pip install datasets)"),
|
||||
raise ImportError(
|
||||
"datasets library is required for HuggingFace mode. "
|
||||
"Install with: pip install datasets, or use local JSON mode instead."
|
||||
)
|
||||
|
||||
load_kw: dict[str, Any] = {
|
||||
@@ -712,10 +626,7 @@ class DailyOmniDataset(BenchmarkDataset):
|
||||
}
|
||||
if self.dataset_subset is not None:
|
||||
load_kw["name"] = self.dataset_subset
|
||||
try:
|
||||
ds = load_dataset(self.dataset_path, **load_kw)
|
||||
except Exception as e:
|
||||
raise _daily_omni_qa_load_error(self.dataset_path, snapshot_error, e) from e
|
||||
ds = load_dataset(self.dataset_path, **load_kw)
|
||||
if not getattr(self, "disable_shuffle", False):
|
||||
ds = ds.shuffle(seed=self.random_seed)
|
||||
self.data = ds
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
Configuration module for vLLM-Omni.
|
||||
"""
|
||||
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config.lora import LoRAConfig
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
from vllm_omni.config.omni_config import (
|
||||
@@ -21,6 +22,7 @@ from vllm_omni.config.omni_config import (
|
||||
VllmOmniGenerationStageConfig,
|
||||
VllmOmniOrchestratorConfig,
|
||||
)
|
||||
from vllm_omni.config.pipeline_registry import register_pipeline
|
||||
from vllm_omni.config.stage_config import (
|
||||
PIPELINE_WIDE_ENGINE_FIELDS,
|
||||
DeployConfig,
|
||||
@@ -40,15 +42,6 @@ from vllm_omni.config.yaml_util import (
|
||||
to_dict,
|
||||
)
|
||||
|
||||
# StageConfigFactory / register_pipeline pull pipeline_registry, which eagerly
|
||||
# imports PI0_PIPELINE → diffusion.data. Keep those lazy so
|
||||
# `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still
|
||||
# loading) cannot close a circular import through DiffusionOutput.
|
||||
_LAZY_ATTRS = {
|
||||
"StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"),
|
||||
"register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"),
|
||||
}
|
||||
|
||||
__all__ = [
|
||||
# Legacy model-level configs.
|
||||
"LoRAConfig",
|
||||
@@ -89,19 +82,3 @@ __all__ = [
|
||||
"merge_configs",
|
||||
"to_dict",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
spec = _LAZY_ATTRS.get(name)
|
||||
if spec is None:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
module_name, attr_name = spec
|
||||
import importlib
|
||||
|
||||
value = getattr(importlib.import_module(module_name), attr_name)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return sorted(set(__all__) | set(globals()))
|
||||
@@ -146,11 +146,6 @@ class OmniModelConfig(ModelConfig):
|
||||
task_type: str | None = None
|
||||
enable_sleep_mode: bool = False
|
||||
has_sampling_extra_args: bool = False
|
||||
# Key names (not values) of the stage's default sampling ``extra_args``.
|
||||
# Engine-core code runs before any request arrives, so this is the only
|
||||
# place it can learn which request-shaping conventions a stage uses (e.g.
|
||||
# ``cfg_role`` for classifier-free-guidance request pairs).
|
||||
sampling_extra_args_keys: tuple[str, ...] = ()
|
||||
|
||||
@property
|
||||
def registry(self):
|
||||
|
||||
+51
-326
@@ -12,22 +12,12 @@ from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import InitVar, dataclass, field, fields
|
||||
from functools import wraps
|
||||
from inspect import Parameter, signature
|
||||
from dataclasses import dataclass, field, fields
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, TypeAlias, TypedDict, cast
|
||||
|
||||
from pydantic import ConfigDict, Field, model_validator
|
||||
from typing_extensions import Self
|
||||
from vllm.config import CacheConfig as VllmCacheConfig
|
||||
from vllm.config import CompilationConfig as VllmCompilationConfig
|
||||
from vllm.config import LoadConfig as VllmLoadConfig
|
||||
from vllm.config import ParallelConfig as VllmParallelConfig
|
||||
from vllm.config import ProfilerConfig as VllmProfilerConfig
|
||||
from vllm.config import SchedulerConfig as VllmSchedulerConfig
|
||||
from pydantic import ConfigDict, Field
|
||||
from vllm.config.utils import config
|
||||
from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
@@ -87,49 +77,6 @@ class _QuantizationEngineOverrides(TypedDict, total=False):
|
||||
quantization: str
|
||||
|
||||
|
||||
class _TrackExplicitConfigFields:
|
||||
"""Record constructor inputs without adding a serialized config field."""
|
||||
|
||||
@model_validator(mode="wrap")
|
||||
@classmethod
|
||||
def _record_explicit_config_fields(cls, value: Any, handler: Any) -> Any:
|
||||
result = handler(value)
|
||||
kwargs = getattr(value, "kwargs", None)
|
||||
if kwargs is not None:
|
||||
explicit_fields = frozenset(kwargs)
|
||||
elif isinstance(value, cls):
|
||||
explicit_fields = getattr(value, "_omni_explicit_fields", frozenset())
|
||||
else:
|
||||
explicit_fields = frozenset()
|
||||
object.__setattr__(result, "_omni_explicit_fields", explicit_fields)
|
||||
return result
|
||||
|
||||
|
||||
def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]:
|
||||
"""Make inherited Pydantic dataclass fields keyword-only as well."""
|
||||
generated_init = cls.__init__
|
||||
generated_signature = signature(cls)
|
||||
keyword_only_signature = generated_signature.replace(
|
||||
parameters=[
|
||||
parameter.replace(kind=Parameter.KEYWORD_ONLY)
|
||||
if parameter.kind in {Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD}
|
||||
else parameter
|
||||
for parameter in generated_signature.parameters.values()
|
||||
]
|
||||
)
|
||||
|
||||
@wraps(generated_init)
|
||||
def keyword_only_init(self: Any, *args: Any, **kwargs: Any) -> None:
|
||||
if args:
|
||||
raise TypeError(f"{cls.__name__}() accepts keyword arguments only; got {len(args)} positional argument(s)")
|
||||
generated_init(self, **kwargs)
|
||||
|
||||
cls.__init__ = keyword_only_init
|
||||
cls.__signature__ = keyword_only_signature
|
||||
cls.__match_args__ = ()
|
||||
return cls
|
||||
|
||||
|
||||
class _ModelEngineOverrides(TypedDict, total=False):
|
||||
model: str
|
||||
model_arch: str
|
||||
@@ -159,6 +106,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
|
||||
enforce_eager: bool
|
||||
max_cudagraph_capture_size: int
|
||||
enable_flashinfer_autotune: bool
|
||||
compilation_config: dict[str, Any]
|
||||
enable_multithread_weight_load: bool
|
||||
num_weight_load_threads: int
|
||||
disable_autocast: bool
|
||||
@@ -200,6 +148,7 @@ class _RuntimeEngineOverrides(TypedDict, total=False):
|
||||
num_gpus: int
|
||||
log_level: str
|
||||
log_stats: bool
|
||||
profiler_config: dict[str, Any]
|
||||
|
||||
|
||||
class _ParallelConfigEngineOverrides(TypedDict, total=False):
|
||||
@@ -243,8 +192,6 @@ class _StageEngineValues:
|
||||
runtime: _RuntimeEngineOverrides
|
||||
parallel: _ParallelEngineOverrides
|
||||
diffusion: _DiffusionEngineOverrides
|
||||
compilation_config: Mapping[str, Any] | VllmCompilationConfig | None
|
||||
profiler_config: Mapping[str, Any] | VllmProfilerConfig | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -394,6 +341,7 @@ class OmniStageModelConfig:
|
||||
enforce_eager: bool = False
|
||||
max_cudagraph_capture_size: int | None = Field(default=None, ge=0)
|
||||
enable_flashinfer_autotune: bool | None = None
|
||||
compilation_config: dict[str, Any] | None = None
|
||||
enable_multithread_weight_load: bool = True
|
||||
num_weight_load_threads: int = Field(default=4, ge=1)
|
||||
disable_autocast: bool = False
|
||||
@@ -404,21 +352,21 @@ class OmniStageModelConfig:
|
||||
tokenizer_subdir: str | None = None
|
||||
|
||||
|
||||
@_enforce_keyword_only_init
|
||||
@config(kw_only=True)
|
||||
class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig):
|
||||
"""vLLM loading behavior plus Omni stage-specific tokenizer inputs."""
|
||||
@config
|
||||
class OmniStageLoadConfig:
|
||||
"""Per-stage loading behavior and resolved tokenizer input."""
|
||||
|
||||
tokenizer: str | None = None
|
||||
download_dir: str | None = None
|
||||
skip_tokenizer_init: bool = False
|
||||
load_format: str = "auto"
|
||||
tokenizer_mode: str = "auto"
|
||||
config_format: str | None = None
|
||||
skip_mm_profiling: bool | None = None
|
||||
|
||||
|
||||
@_enforce_keyword_only_init
|
||||
@config(kw_only=True)
|
||||
class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig):
|
||||
@config
|
||||
class OmniStageCacheConfig:
|
||||
"""Per-stage engine cache and memory behavior.
|
||||
|
||||
This is separate from ``_DiffusionConfigProjection.cache_config``, which configures
|
||||
@@ -435,26 +383,17 @@ class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig):
|
||||
mamba_ssm_cache_dtype: str | None = None
|
||||
|
||||
|
||||
@_enforce_keyword_only_init
|
||||
@config(kw_only=True)
|
||||
class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig):
|
||||
@config
|
||||
class OmniStageSchedulerConfig:
|
||||
"""Per-stage request scheduling behavior."""
|
||||
|
||||
# Upstream receives max_model_len only while materializing SchedulerConfig.
|
||||
# Omni retains it as unresolved stage input until the owning engine process.
|
||||
max_num_seqs: int | None = Field(default=None, ge=1)
|
||||
max_num_batched_tokens: int | None = Field(default=None, ge=1)
|
||||
max_model_len: int | None = Field(default=None, ge=-1)
|
||||
is_encoder_decoder: InitVar[bool] = False
|
||||
enable_chunked_prefill: bool | None = None
|
||||
async_scheduling: bool | None = None
|
||||
|
||||
def __post_init__(self, is_encoder_decoder: bool = False) -> None:
|
||||
# Upstream initializes these derived fields in its terminal post-init.
|
||||
# Keep them serializable here without running model-dependent checks.
|
||||
self.max_num_encoder_input_tokens = self.max_num_batched_tokens
|
||||
self.encoder_cache_size = self.max_num_batched_tokens
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if (
|
||||
self.max_num_batched_tokens is not None
|
||||
and self.max_num_seqs is not None
|
||||
@@ -493,73 +432,24 @@ class OmniStageRuntimeConfig:
|
||||
num_gpus: int = Field(default=1, ge=1)
|
||||
log_level: str = "info"
|
||||
log_stats: bool = False
|
||||
profiler_config: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@_enforce_keyword_only_init
|
||||
@config(kw_only=True)
|
||||
class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig):
|
||||
@config
|
||||
class OmniStageParallelConfig:
|
||||
"""Common per-stage distributed parallelism behavior."""
|
||||
|
||||
# EngineArgs intentionally leaves these unresolved. The upstream terminal
|
||||
# ParallelConfig defaults (rank 0, local size 1, port 29550, worker
|
||||
# ``"auto"``) must not turn into explicit engine inputs in the head
|
||||
# process merely because this transport class inherits ParallelConfig.
|
||||
data_parallel_size_local: int | None = Field(default=None, ge=0)
|
||||
data_parallel_rank: int | None = Field(default=None, ge=0)
|
||||
data_parallel_rpc_port: int | None = None
|
||||
worker_cls: str | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_parallel_config(self) -> Self:
|
||||
"""Run upstream validation without resolving deferred engine inputs."""
|
||||
deferred_values = (
|
||||
self.data_parallel_size_local,
|
||||
self.data_parallel_rank,
|
||||
self.data_parallel_rpc_port,
|
||||
self.worker_cls,
|
||||
self.all2all_backend,
|
||||
self.disable_custom_all_reduce,
|
||||
)
|
||||
if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size:
|
||||
raise ValueError(
|
||||
f"data_parallel_size_local ({self.data_parallel_size_local}) "
|
||||
f"must be <= data_parallel_size ({self.data_parallel_size})"
|
||||
)
|
||||
if self.data_parallel_rank is not None and not 0 <= self.data_parallel_rank < self.data_parallel_size:
|
||||
raise ValueError(
|
||||
f"data_parallel_rank ({self.data_parallel_rank}) must be in the range [0, {self.data_parallel_size})"
|
||||
)
|
||||
self.data_parallel_size_local = self.data_parallel_size
|
||||
self.data_parallel_rank = 0
|
||||
self.data_parallel_rpc_port = VllmParallelConfig.data_parallel_rpc_port
|
||||
self.worker_cls = VllmParallelConfig.worker_cls
|
||||
try:
|
||||
VllmParallelConfig._validate_parallel_config(self)
|
||||
finally:
|
||||
(
|
||||
self.data_parallel_size_local,
|
||||
self.data_parallel_rank,
|
||||
self.data_parallel_rpc_port,
|
||||
self.worker_cls,
|
||||
self.all2all_backend,
|
||||
self.disable_custom_all_reduce,
|
||||
) = deferred_values
|
||||
return self
|
||||
pipeline_parallel_size: int = Field(default=1, ge=1)
|
||||
data_parallel_size: int = Field(default=1, ge=1)
|
||||
tensor_parallel_size: int = Field(default=1, ge=1)
|
||||
enable_expert_parallel: bool = False
|
||||
world_size: int = Field(default=1, ge=1, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Keep config construction transport-safe. Upstream runtime backend,
|
||||
# rank, port, and platform resolution happens in the owning process.
|
||||
self.data_parallel_index = self.data_parallel_rank
|
||||
self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size
|
||||
|
||||
@property
|
||||
def world_size_across_dp(self) -> int:
|
||||
# Omni's public world_size has historically included DP.
|
||||
return self.world_size
|
||||
|
||||
|
||||
@_enforce_keyword_only_init
|
||||
@config(kw_only=True)
|
||||
@config
|
||||
class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
|
||||
"""Diffusion-stage distributed parallelism behavior."""
|
||||
|
||||
@@ -578,7 +468,6 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
|
||||
hsdp_replicate_size: int = Field(default=1, ge=1)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.data_parallel_index = self.data_parallel_rank
|
||||
self.sequence_parallel_size = (
|
||||
self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
|
||||
)
|
||||
@@ -600,17 +489,12 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
|
||||
* self.cfg_parallel_size
|
||||
)
|
||||
if self.use_hsdp:
|
||||
incompatible = []
|
||||
if self.tensor_parallel_size > 1:
|
||||
incompatible.append("TP")
|
||||
if self.data_parallel_size > 1:
|
||||
incompatible.append("DP")
|
||||
if self.pipeline_parallel_size > 1:
|
||||
incompatible.append("PP")
|
||||
if self.enable_expert_parallel:
|
||||
incompatible.append("EP")
|
||||
if incompatible:
|
||||
raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
|
||||
if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
|
||||
raise ValueError(
|
||||
"HSDP (use_hsdp=True) cannot be used with TP or DP "
|
||||
f"(tensor_parallel_size={self.tensor_parallel_size}, "
|
||||
f"data_parallel_size={self.data_parallel_size})"
|
||||
)
|
||||
if self.hsdp_shard_size == -1:
|
||||
if other_parallel_world_size == 1:
|
||||
raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1")
|
||||
@@ -930,10 +814,6 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
|
||||
"omni_kv_config",
|
||||
"cfg_kv_collect_func",
|
||||
"max_num_seqs",
|
||||
"kv_cache_memory_bytes",
|
||||
"gpu_memory_utilization",
|
||||
"max_num_batched_tokens",
|
||||
"max_model_len",
|
||||
"enable_sleep_mode",
|
||||
"enforce_eager",
|
||||
"enable_multithread_weight_load",
|
||||
@@ -959,58 +839,6 @@ _DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMP
|
||||
"stage_id",
|
||||
}
|
||||
|
||||
|
||||
def _upstream_engine_field_map(
|
||||
config_cls: type[Any],
|
||||
*,
|
||||
aliases: Mapping[str, str] = {},
|
||||
exclude: frozenset[str] = frozenset(),
|
||||
) -> dict[str, str]:
|
||||
"""Map reusable upstream config fields to their EngineArgs inputs."""
|
||||
engine_fields = frozenset(config_field.name for config_field in fields(VllmEngineArgs))
|
||||
return {
|
||||
config_field.name: engine_name
|
||||
for config_field in fields(config_cls)
|
||||
if config_field.init
|
||||
and config_field.name not in exclude
|
||||
and (engine_name := aliases.get(config_field.name, config_field.name)) in engine_fields
|
||||
}
|
||||
|
||||
|
||||
_LOAD_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(VllmLoadConfig)
|
||||
_CACHE_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
|
||||
VllmCacheConfig,
|
||||
aliases={"cache_dtype": "kv_cache_dtype"},
|
||||
)
|
||||
_SCHEDULER_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
|
||||
VllmSchedulerConfig,
|
||||
aliases={"policy": "scheduling_policy"},
|
||||
# ``scheduler_cls`` is selected by the immutable stage topology, while
|
||||
# ``disable_hybrid_kv_cache_manager`` belongs to the cache concern in the
|
||||
# Omni schema. Keep both out of the dynamic SchedulerConfig projection so
|
||||
# one input cannot acquire two owners.
|
||||
exclude=frozenset(
|
||||
{
|
||||
"scheduler_cls",
|
||||
"disable_hybrid_kv_cache_manager",
|
||||
}
|
||||
),
|
||||
)
|
||||
_PARALLEL_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
|
||||
VllmParallelConfig,
|
||||
aliases={"data_parallel_master_ip": "data_parallel_address"},
|
||||
# Runtime owns worker selection. The private API-process fields are
|
||||
# terminal vLLM internals rather than per-stage user inputs.
|
||||
exclude=frozenset(
|
||||
{
|
||||
"distributed_executor_backend",
|
||||
"worker_cls",
|
||||
"_api_process_count",
|
||||
"_api_process_rank",
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
_QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__)
|
||||
_MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
|
||||
_LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
|
||||
@@ -1018,17 +846,7 @@ _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
|
||||
_SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
|
||||
_CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
|
||||
_RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
|
||||
_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"})
|
||||
_LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values())
|
||||
_LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values())
|
||||
_LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values())
|
||||
_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP.values())
|
||||
_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(
|
||||
f.name for f in fields(OmniStageDiffusionParallelConfig)
|
||||
) & frozenset(_ParallelConfigEngineOverrides.__annotations__)
|
||||
_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP = {name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS}
|
||||
_LLM_PARALLEL_CONFIG_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP)
|
||||
_PARALLEL_CONFIG_ENGINE_FIELDS = _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
|
||||
_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__)
|
||||
_PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
|
||||
_COMMON_STAGE_ENGINE_FIELDS = (
|
||||
_QUANTIZATION_ENGINE_FIELDS
|
||||
@@ -1038,16 +856,12 @@ _COMMON_STAGE_ENGINE_FIELDS = (
|
||||
| _SCHEDULER_ENGINE_FIELDS
|
||||
| _CONNECTOR_ENGINE_FIELDS
|
||||
| _RUNTIME_ENGINE_FIELDS
|
||||
| _DIRECT_VLLM_CONFIG_ENGINE_FIELDS
|
||||
)
|
||||
_LLM_STAGE_ENGINE_FIELDS = (
|
||||
_COMMON_STAGE_ENGINE_FIELDS
|
||||
| _LLM_LOAD_ENGINE_FIELDS
|
||||
| _LLM_CACHE_ENGINE_FIELDS
|
||||
| _LLM_SCHEDULER_ENGINE_FIELDS
|
||||
| _LLM_PARALLEL_CONFIG_ENGINE_FIELDS
|
||||
| {"parallel_config"}
|
||||
)
|
||||
_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"}
|
||||
_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - {
|
||||
"world_size"
|
||||
}
|
||||
_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
|
||||
_DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
|
||||
_COMMON_STAGE_ENGINE_FIELDS
|
||||
| _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
|
||||
@@ -1068,27 +882,6 @@ _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
|
||||
StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
|
||||
StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS,
|
||||
}
|
||||
_PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = {
|
||||
StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS,
|
||||
StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS,
|
||||
StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP),
|
||||
}
|
||||
|
||||
_LOAD_STAGE_ENGINE_FIELD_MAP = {
|
||||
**{name: name for name in _LOAD_ENGINE_FIELDS},
|
||||
**_LOAD_CONFIG_ENGINE_FIELD_MAP,
|
||||
}
|
||||
_CACHE_STAGE_ENGINE_FIELD_MAP = {
|
||||
**{name: name for name in _CACHE_ENGINE_FIELDS},
|
||||
**_CACHE_CONFIG_ENGINE_FIELD_MAP,
|
||||
}
|
||||
_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {
|
||||
**{name: name for name in _SCHEDULER_ENGINE_FIELDS},
|
||||
**_SCHEDULER_CONFIG_ENGINE_FIELD_MAP,
|
||||
}
|
||||
_DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP = {name: name for name in _LOAD_ENGINE_FIELDS}
|
||||
_DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP = {name: name for name in _CACHE_ENGINE_FIELDS}
|
||||
_DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {name: name for name in _SCHEDULER_ENGINE_FIELDS}
|
||||
|
||||
|
||||
def _validate_stage_engine_override_ownership(
|
||||
@@ -1099,19 +892,13 @@ def _validate_stage_engine_override_ownership(
|
||||
try:
|
||||
owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
|
||||
parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
|
||||
parallel_config_fields = _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE[execution_type]
|
||||
except KeyError as exc:
|
||||
raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc
|
||||
|
||||
unowned_fields = set(overrides) - owner_fields
|
||||
parallel_config = overrides.get("parallel_config")
|
||||
if isinstance(parallel_config, Mapping):
|
||||
# Nested values traditionally use upstream config names, while flat
|
||||
# EngineArgs uses aliases such as ``data_parallel_address``. Accept
|
||||
# either spelling at the boundary; the builder canonicalizes aliases
|
||||
# before instantiating the inherited config.
|
||||
accepted_parallel_names = parallel_config_fields | parallel_owner_fields
|
||||
unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - accepted_parallel_names)
|
||||
unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields)
|
||||
if unowned_fields:
|
||||
names = ", ".join(sorted(unowned_fields))
|
||||
raise ValueError(
|
||||
@@ -1182,28 +969,17 @@ def _stage_engine_values(
|
||||
topology.execution_type,
|
||||
engine,
|
||||
)
|
||||
if topology.execution_type in {
|
||||
StageExecutionType.LLM_AR,
|
||||
StageExecutionType.LLM_GENERATION,
|
||||
}:
|
||||
load_engine_fields = _LLM_LOAD_ENGINE_FIELDS
|
||||
cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS
|
||||
scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS
|
||||
else:
|
||||
load_engine_fields = _LOAD_ENGINE_FIELDS
|
||||
cache_engine_fields = _CACHE_ENGINE_FIELDS
|
||||
scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS
|
||||
return _StageEngineValues(
|
||||
quantization=cast(
|
||||
_QuantizationEngineOverrides,
|
||||
_select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS),
|
||||
),
|
||||
model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)),
|
||||
load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)),
|
||||
cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)),
|
||||
load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)),
|
||||
cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)),
|
||||
scheduler=cast(
|
||||
_SchedulerEngineOverrides,
|
||||
_select_engine_overrides(engine, scheduler_engine_fields),
|
||||
_select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS),
|
||||
),
|
||||
connector=cast(
|
||||
_ConnectorEngineOverrides,
|
||||
@@ -1212,8 +988,6 @@ def _stage_engine_values(
|
||||
runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)),
|
||||
parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)),
|
||||
diffusion=_DiffusionEngineOverrides.from_engine(engine),
|
||||
compilation_config=_copy_value(engine.get("compilation_config")),
|
||||
profiler_config=_copy_value(engine.get("profiler_config")),
|
||||
)
|
||||
|
||||
|
||||
@@ -1268,8 +1042,6 @@ class BaseVllmOmniStageConfig:
|
||||
connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig)
|
||||
runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig)
|
||||
parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig)
|
||||
compilation_config: VllmCompilationConfig | None = None
|
||||
profiler_config: VllmProfilerConfig | None = None
|
||||
quantization_config: _QuantizationConfigType = None
|
||||
|
||||
@property
|
||||
@@ -1421,8 +1193,6 @@ def _build_common_stage_config_kwargs(
|
||||
parallel_config,
|
||||
),
|
||||
"parallel_config": parallel_config,
|
||||
"compilation_config": _copy_value(engine.compilation_config),
|
||||
"profiler_config": _copy_value(engine.profiler_config),
|
||||
"quantization_config": _copy_value(quantization_config),
|
||||
},
|
||||
input_proc,
|
||||
@@ -1617,51 +1387,12 @@ def _build_load_config(
|
||||
return OmniStageLoadConfig(**kwargs)
|
||||
|
||||
|
||||
def _config_kwargs_from_engine_args(
|
||||
engine: Mapping[str, Any],
|
||||
field_map: Mapping[str, str],
|
||||
) -> dict[str, Any]:
|
||||
"""Convert EngineArgs names back to their upstream config field names."""
|
||||
return {
|
||||
config_name: _copy_value(engine[engine_name])
|
||||
for config_name, engine_name in field_map.items()
|
||||
if engine_name in engine and engine[engine_name] is not None
|
||||
}
|
||||
|
||||
|
||||
def _normalize_config_mapping(
|
||||
values: Mapping[str, Any],
|
||||
field_map: Mapping[str, str],
|
||||
) -> dict[str, Any]:
|
||||
"""Canonicalize config and EngineArgs spellings to config field names."""
|
||||
engine_to_config = {engine_name: config_name for config_name, engine_name in field_map.items()}
|
||||
normalized: dict[str, Any] = {}
|
||||
source_names: dict[str, str] = {}
|
||||
for name, value in values.items():
|
||||
if value is None:
|
||||
continue
|
||||
config_name = name if name in field_map else engine_to_config.get(name, name)
|
||||
if config_name in normalized:
|
||||
raise ValueError(
|
||||
f"Config mapping specifies both {source_names[config_name]!r} and {name!r} "
|
||||
f"for upstream field {config_name!r}"
|
||||
)
|
||||
normalized[config_name] = _copy_value(value)
|
||||
source_names[config_name] = name
|
||||
return normalized
|
||||
|
||||
|
||||
def _omni_config_kwargs(engine: Mapping[str, Any], fields_: frozenset[str]) -> dict[str, Any]:
|
||||
return {name: _copy_value(engine[name]) for name in fields_ if name in engine and engine[name] is not None}
|
||||
|
||||
|
||||
def _build_cache_config(
|
||||
deploy: DeployConfig,
|
||||
engine: _CacheEngineOverrides,
|
||||
execution_type: StageExecutionType,
|
||||
) -> OmniStageCacheConfig:
|
||||
kwargs = _config_kwargs_from_engine_args(engine, _CACHE_CONFIG_ENGINE_FIELD_MAP)
|
||||
kwargs.update(_omni_config_kwargs(engine, _CACHE_ENGINE_FIELDS))
|
||||
kwargs = _config_kwargs(engine)
|
||||
if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None:
|
||||
kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching)
|
||||
if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION:
|
||||
@@ -1675,8 +1406,7 @@ def _build_scheduler_config(
|
||||
engine: _SchedulerEngineOverrides,
|
||||
execution_type: StageExecutionType,
|
||||
) -> OmniStageSchedulerConfig:
|
||||
kwargs = _config_kwargs_from_engine_args(engine, _SCHEDULER_CONFIG_ENGINE_FIELD_MAP)
|
||||
kwargs.update(_omni_config_kwargs(engine, _SCHEDULER_ENGINE_FIELDS))
|
||||
kwargs = _config_kwargs(engine)
|
||||
if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None:
|
||||
kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill)
|
||||
if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR:
|
||||
@@ -1724,21 +1454,16 @@ def _build_parallel_config(
|
||||
config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig,
|
||||
) -> OmniStageParallelConfig:
|
||||
parallel_config = _mapping_or_empty(engine.get("parallel_config"))
|
||||
config_fields = frozenset(signature(config_cls).parameters)
|
||||
is_diffusion = issubclass(config_cls, OmniStageDiffusionParallelConfig)
|
||||
field_map = _DIFFUSION_PARALLEL_CONFIG_FIELD_MAP if is_diffusion else _PARALLEL_CONFIG_ENGINE_FIELD_MAP
|
||||
owned_parallel_fields = frozenset(field_map)
|
||||
kwargs = _config_kwargs(engine) if is_diffusion else _config_kwargs_from_engine_args(engine, field_map)
|
||||
kwargs = {name: value for name, value in kwargs.items() if name in owned_parallel_fields}
|
||||
config_fields = {
|
||||
config_field.name
|
||||
for config_field in fields(config_cls)
|
||||
if getattr(config_field.default, "init", None) is not False
|
||||
}
|
||||
kwargs = {
|
||||
name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None
|
||||
}
|
||||
kwargs = {name: value for name, value in kwargs.items() if name in config_fields}
|
||||
nested_kwargs = _normalize_config_mapping(parallel_config, field_map)
|
||||
kwargs.update(
|
||||
{
|
||||
name: value
|
||||
for name, value in nested_kwargs.items()
|
||||
if name in owned_parallel_fields and name in config_fields
|
||||
}
|
||||
)
|
||||
kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields})
|
||||
if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None:
|
||||
kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size)
|
||||
if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None:
|
||||
|
||||
@@ -33,7 +33,7 @@ from vllm.logger import init_logger
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
|
||||
from vllm_omni.diffusion.models.pi0.pipeline_pi0 import PI0_PIPELINE
|
||||
from vllm_omni.model_executor.models.audex.pipeline import (
|
||||
AUDEX_S2S_PIPELINE,
|
||||
AUDEX_THINKER_ONLY_PIPELINE,
|
||||
@@ -85,7 +85,6 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import (
|
||||
MING_TTS_PIPELINE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
|
||||
from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE
|
||||
from vllm_omni.model_executor.models.moss_tts.pipeline import (
|
||||
MOSS_TTS_LOCAL_PIPELINE,
|
||||
MOSS_TTS_PIPELINE,
|
||||
@@ -178,7 +177,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
|
||||
"moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
|
||||
"minicpmo_4_5": MINICPMO_4_5_PIPELINE,
|
||||
"minimax_music3": MINIMAX_MUSIC3_PIPELINE,
|
||||
"higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
|
||||
"higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,
|
||||
"dynin_omni": DYNIN_OMNI_PIPELINE,
|
||||
|
||||
@@ -91,11 +91,6 @@ class OmniPayloadMeta(TypedDict, total=False):
|
||||
ref_context_included: bool
|
||||
talker_prefill_offset: int
|
||||
omni_final_stage_id: int
|
||||
# Per-request audio seed. Stages that draw their own noise (flow-matching
|
||||
# / diffusion decoders) need the producing stage's seed to stay
|
||||
# reproducible, and the producing stage's SamplingParams do not travel
|
||||
# with the payload.
|
||||
audio_seed: int
|
||||
|
||||
|
||||
class OmniPayload(TypedDict, total=False):
|
||||
@@ -194,7 +189,6 @@ class MetaStruct(_StructBase):
|
||||
codec_left_context_frames: int | None = None
|
||||
code_flat_numel: int | None = None
|
||||
omni_final_stage_id: int | None = None
|
||||
audio_seed: int | None = None
|
||||
|
||||
|
||||
class OmniPayloadStruct(_StructBase):
|
||||
|
||||
@@ -1,130 +0,0 @@
|
||||
# MiniMax-Music3 deploy: Stage 0 (AR talker) -> Stage 1 (acoustic decoder),
|
||||
# both on one 140 GB card, streaming 200-frame windows over shared memory.
|
||||
#
|
||||
# Auto-discovered default for model_type=minimax_music3.
|
||||
#
|
||||
# Classifier-free guidance is mandatory for this checkpoint: every request
|
||||
# decodes as a cond/uncond pair in stage 0, so stage-0 KV, max_num_seqs and
|
||||
# prefill budget are all sized for TWO rows per user request. The stage-0
|
||||
# extra_args block below is what turns on has_sampling_extra_args (so per-row
|
||||
# sampling extra args reach the model's forward) and is also how the CFG
|
||||
# scheduler patches learn this stage is paired.
|
||||
pipeline: minimax_music3
|
||||
async_chunk: true
|
||||
|
||||
# The repo-root config.json reports model_type=minimax_music3, which upstream
|
||||
# transformers does not ship a config class for.
|
||||
trust_remote_code: true
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
# Stream stereo PCM back to the client as stage 1 finishes each window.
|
||||
codec_streaming: true
|
||||
# Window geometry is fixed by the checkpoint (200-frame windows on a
|
||||
# 100-frame hop, models/minimax_music3/constants.py), so there are no
|
||||
# codec_chunk_frames knobs here: changing the window changes the audio.
|
||||
shm_threshold_bytes: 65536
|
||||
connector_get_sleep_s: 0.01
|
||||
# One window is 200 AR frames = 8 s of music, so stage 1 waits far longer
|
||||
# for a chunk than a per-step codec stage would. First chunk additionally
|
||||
# covers stage-0 prefill of both CFG rows.
|
||||
connector_get_max_wait_first_chunk: 30000
|
||||
connector_get_max_wait: 10000
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
# 16 rows = 8 concurrent guided requests, since CFG doubles every request.
|
||||
max_num_seqs: 16
|
||||
# ~16 GB of bf16 Qwen3-8B weights plus KV for 16 full-length rows
|
||||
# (36 layers x 8 KV heads x 128 head_dim ~= 144 KB/token, so ~1.4 GB per
|
||||
# 10240-token row) fits comfortably inside 0.6 of a 140 GB card.
|
||||
gpu_memory_utilization: 0.6
|
||||
dtype: bfloat16
|
||||
# The checkpoint's max_position_embeddings is 10240 and the whole song is
|
||||
# decoded in-context. Note this cannot hold a max-length prompt AND a
|
||||
# max-length song at the same time: the prompt cap is 5000 tokens and the
|
||||
# frame budget is 9000, which overruns 10240. That is a property of the
|
||||
# checkpoint, not of this file; long captions shorten the maximum song.
|
||||
max_model_len: 10240
|
||||
# One scheduler step must be able to prefill both rows of a pair
|
||||
# (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked.
|
||||
max_num_batched_tokens: 10240
|
||||
# Two rows at 8B is a launch-bound decode step, so the backbone is
|
||||
# captured. FULL_DECODE_ONLY only ever dispatches a graph for a uniform
|
||||
# decode batch; prefill and mixed steps stay eager, which is what the
|
||||
# talker's frame-feedback fallback is for. The talker's own sampler and
|
||||
# depth decoder run after the graph, not inside it.
|
||||
#
|
||||
# What makes the forward capturable: the frame feedback writes a
|
||||
# fixed-size slab in prepare_runner_inputs (outside the captured region)
|
||||
# and the forward selects over it with torch.where, instead of an
|
||||
# index_copy against an index tensor whose address and length changed
|
||||
# every step. The row-to-request binding the sampler reads is likewise
|
||||
# rebuilt in prepare_runner_inputs, because a replayed step never runs
|
||||
# the forward's Python body.
|
||||
enforce_eager: false
|
||||
compilation_config:
|
||||
# No compilation. The talker walks the backbone's decoder layers itself,
|
||||
# so the compiled Qwen3Model.forward is never entered and inductor has
|
||||
# nothing to do; what the mode still decides is which kernels vLLM
|
||||
# defaults to. Left at its non-eager default it picks the elementwise
|
||||
# PyTorch norms over the fused vllm_c ones, which is both slower and
|
||||
# numerically different: measured here, that alone moved the sampled
|
||||
# codes off the eager trajectory by frame 4. Graph capture is
|
||||
# independent of this and still happens.
|
||||
mode: 0
|
||||
# 16 rows is the stage ceiling, i.e. 8 guided requests.
|
||||
cudagraph_capture_sizes: [1, 2, 4, 8, 16]
|
||||
cudagraph_mode: FULL_DECODE_ONLY
|
||||
cudagraph_num_of_warmups: 1
|
||||
# Belt and braces on the same point: the non-eager default is "none",
|
||||
# which hands the norm and activation ops to elementwise PyTorch.
|
||||
custom_ops: ["all"]
|
||||
# A CFG pair must be sampled inside one engine step for the logits blend
|
||||
# and the cond->uncond token copy to line up.
|
||||
async_scheduling: false
|
||||
# Off during bring-up: the uncond row is a long run of identical
|
||||
# <|audio_cfg|> tokens and would take a much larger prefix-cache hit than
|
||||
# its cond partner, which is exactly the num_computed_tokens skew the CFG
|
||||
# pair patches then have to undo.
|
||||
enable_prefix_caching: false
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
# The frame budget is enforced by the model. This ceiling only has to
|
||||
# leave room for the extra step that delivers the final window.
|
||||
max_tokens: 9002
|
||||
extra_args:
|
||||
cfg_role: cond
|
||||
cfg_scale: 1.5
|
||||
|
||||
- stage_id: 1
|
||||
# One row per user request: the uncond companion is stage-0-final and
|
||||
# never reaches the acoustic stage.
|
||||
max_num_seqs: 8
|
||||
# No KV cache on this stage; the budget only covers the DiT and DAV
|
||||
# weights plus per-window latents and activations.
|
||||
gpu_memory_utilization: 0.2
|
||||
# The flow-matching solver degrades measurably in bfloat16
|
||||
# (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py).
|
||||
dtype: float32
|
||||
enforce_eager: true
|
||||
# Async scheduling builds per-request payloads against the NEXT step's
|
||||
# input batch, which breaks when chunk arrival staggers requests.
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
# One placeholder prompt token per arriving window, so a six-minute song
|
||||
# (90 windows) needs a tiny context.
|
||||
max_model_len: 4096
|
||||
max_num_batched_tokens: 4096
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
max_tokens: 4096
|
||||
seed: 0
|
||||
@@ -1,97 +0,0 @@
|
||||
# MiniMax-Music3 two-GPU layout: AR talker on device 0, acoustic decoder on
|
||||
# device 1.
|
||||
#
|
||||
# Splitting the stages removes the memory-budget compromise of the single-card
|
||||
# file: the talker's CFG-doubled KV no longer competes with the DiT and DAV
|
||||
# weights, and the flow-matching solver no longer competes with AR decode for
|
||||
# SMs. Everything else (window geometry, CFG contract, sampling defaults)
|
||||
# matches minimax_music3.yaml, so the two files stay behaviorally identical.
|
||||
pipeline: minimax_music3
|
||||
async_chunk: true
|
||||
|
||||
# The repo-root config.json reports model_type=minimax_music3, which upstream
|
||||
# transformers does not ship a config class for.
|
||||
trust_remote_code: true
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
# Stream stereo PCM back to the client as stage 1 finishes each window.
|
||||
codec_streaming: true
|
||||
# Window geometry is fixed by the checkpoint (200-frame windows on a
|
||||
# 100-frame hop, models/minimax_music3/constants.py), so there are no
|
||||
# codec_chunk_frames knobs here: changing the window changes the audio.
|
||||
shm_threshold_bytes: 65536
|
||||
connector_get_sleep_s: 0.01
|
||||
# One window is 200 AR frames = 8 s of music, so stage 1 waits far longer
|
||||
# for a chunk than a per-step codec stage would. First chunk additionally
|
||||
# covers stage-0 prefill of both CFG rows.
|
||||
connector_get_max_wait_first_chunk: 30000
|
||||
connector_get_max_wait: 10000
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
# 32 rows = 16 concurrent guided requests, since CFG doubles every request.
|
||||
# Twice the single-card figure: this stage now owns its whole device.
|
||||
max_num_seqs: 32
|
||||
gpu_memory_utilization: 0.85
|
||||
dtype: bfloat16
|
||||
# The checkpoint's max_position_embeddings is 10240 and the whole song is
|
||||
# decoded in-context. Note this cannot hold a max-length prompt AND a
|
||||
# max-length song at the same time: the prompt cap is 5000 tokens and the
|
||||
# frame budget is 9000, which overruns 10240. That is a property of the
|
||||
# checkpoint, not of this file; long captions shorten the maximum song.
|
||||
max_model_len: 10240
|
||||
# One scheduler step must be able to prefill both rows of a pair
|
||||
# (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked.
|
||||
max_num_batched_tokens: 10240
|
||||
# The talker owns its sampler and drives a depth decoder per AR step;
|
||||
# leave graph capture off until that path is verified graph-safe.
|
||||
enforce_eager: true
|
||||
# A CFG pair must be sampled inside one engine step for the logits blend
|
||||
# and the cond->uncond token copy to line up.
|
||||
async_scheduling: false
|
||||
# Off during bring-up: the uncond row is a long run of identical
|
||||
# <|audio_cfg|> tokens and would take a much larger prefix-cache hit than
|
||||
# its cond partner, which is exactly the num_computed_tokens skew the CFG
|
||||
# pair patches then have to undo.
|
||||
enable_prefix_caching: false
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
# The frame budget is enforced by the model. This ceiling only has to
|
||||
# leave room for the extra step that delivers the final window.
|
||||
max_tokens: 9002
|
||||
extra_args:
|
||||
cfg_role: cond
|
||||
cfg_scale: 1.5
|
||||
|
||||
- stage_id: 1
|
||||
# One row per user request: the uncond companion is stage-0-final and
|
||||
# never reaches the acoustic stage.
|
||||
max_num_seqs: 16
|
||||
# No KV cache on this stage, but it now has a device to itself, so the
|
||||
# solver can hold more windows in flight than the shared-card profile.
|
||||
gpu_memory_utilization: 0.6
|
||||
# The flow-matching solver degrades measurably in bfloat16
|
||||
# (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py).
|
||||
dtype: float32
|
||||
enforce_eager: true
|
||||
# Async scheduling builds per-request payloads against the NEXT step's
|
||||
# input batch, which breaks when chunk arrival staggers requests.
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
# One placeholder prompt token per arriving window, so a six-minute song
|
||||
# (90 windows) needs a tiny context.
|
||||
max_model_len: 4096
|
||||
max_num_batched_tokens: 4096
|
||||
devices: "1"
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
max_tokens: 4096
|
||||
seed: 0
|
||||
@@ -81,19 +81,6 @@ class AttentionBackend(ABC):
|
||||
supported_head_sizes = cls.get_supported_head_sizes()
|
||||
return (not supported_head_sizes) or head_size in supported_head_sizes
|
||||
|
||||
@classmethod
|
||||
def indexes_kv_by_block_stride(cls) -> bool:
|
||||
"""Whether this backend reads K/V pages by the runtime block stride.
|
||||
|
||||
Returning ``True`` means the physical cache layout has ``num_blocks``
|
||||
as its outer stride, so native vLLM may safely use page-size padding
|
||||
when it unifies cache layouts across layers. Dense diffusion backends
|
||||
conservatively keep the default ``False``; a paged backend should
|
||||
override this only when its kernel actually follows that layout.
|
||||
"""
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class QueryRange:
|
||||
|
||||
@@ -11,10 +11,8 @@ from dataclasses import replace
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.utils import extract_layer_index
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
|
||||
@@ -63,23 +61,12 @@ class Attention(nn.Module):
|
||||
# perf for this layer (e.g. Wan2.2 cross-attn has short sequences and
|
||||
# block-FP8 quant offers no win). Default False = follow global config.
|
||||
disable_kv_quant: bool = False,
|
||||
# Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion
|
||||
# attention remains dense and contributes no native KVCacheSpec.
|
||||
paged_kv_cache_role: str | None = None,
|
||||
paged_kv_cache_dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.role = role
|
||||
self.role_category = role_category
|
||||
self.qkv_layout = qkv_layout
|
||||
if paged_kv_cache_role == "":
|
||||
raise ValueError("paged_kv_cache_role must be non-empty when provided")
|
||||
self.paged_kv_cache_role = paged_kv_cache_role
|
||||
self.paged_kv_cache_dtype = paged_kv_cache_dtype
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
|
||||
self.head_size = head_size
|
||||
|
||||
# Resolve backend via role-aware config.
|
||||
# The global diffusion config is set during model init via
|
||||
@@ -184,25 +171,6 @@ class Attention(nn.Module):
|
||||
self._disable_kv_quant: bool = disable_kv_quant
|
||||
self._init_kv_cache_quantization(config)
|
||||
|
||||
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
|
||||
"""Return native rank-local geometry for an opted-in paged cache."""
|
||||
|
||||
if self.paged_kv_cache_role is None:
|
||||
return None
|
||||
dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
|
||||
# Keep backend layout discovery under the same config context used by
|
||||
# upstream vLLM's attention-spec collector.
|
||||
with set_current_vllm_config(vllm_config):
|
||||
indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride()
|
||||
return FullAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
head_size=self.head_size,
|
||||
dtype=dtype,
|
||||
indexes_kv_by_block_stride=indexes_kv_by_block_stride,
|
||||
non_causal=not self.causal,
|
||||
)
|
||||
|
||||
def _get_active_parallel_strategy(self):
|
||||
"""Get the parallel strategy based on current SP active state.
|
||||
|
||||
|
||||
Loaded 100 of 199 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user