diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 18bd82a9d..9be72d3d5 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -460,6 +460,21 @@ steps: exit $$EXIT2 mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3" + key: nightly-diffusion-x2v-performance-minimax-h3 + timeout_in_minutes: 180 + commands: + - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results + - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + exit $$EXIT + mirror_hardwares: h100_4 + - group: ":card_index_dividers: Diffusion Test" key: nightly-diffusion-test-group depends_on: upload-nightly-pipeline diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index 63382e76f..bae8cd1b5 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -309,7 +309,7 @@ For Omni or custom models, create: Required updates: 1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table -2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table) +2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table) 3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported 4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported 5. Offline example docs under `examples/offline_inference//` (`README.md` or category-specific `.md`) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 6b3ddb474..c95d1df08 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -16,15 +16,15 @@ /.claude/ @hsliuustc0106 # Tests -/tests/ @yenuo26 +/tests/ @yenuo26 @NickCao # CI -/.buildkite/ @yenuo26 @congw729 +/.buildkite/ @yenuo26 @congw729 @NickCao # Engine and model_executor /vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 /vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307 -/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks +/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao /vllm_omni/model_executor/ @tzhouam @gcanlin /vllm_omni/worker/ @tzhouam @gcanlin @@ -37,7 +37,7 @@ # Omni and TTS support /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao -/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin +/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao # Benchmarks /benchmarks/accuracy/ @david6666666 diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index f8ad1a71e..aeff3ebbb 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -1,5 +1,6 @@ import asyncio import base64 +import json import mimetypes import os import time @@ -13,6 +14,17 @@ from tqdm import tqdm DEFAULT_EDITS_BOT_TASK = "think" +def _is_minimax_h3_mixed_reference(model: str) -> bool: + """MiniMax-H3 Ref2VA consumes reference videos as file paths. + + Its reference preparer ffprobes / ffmpegs the reference file, so single + ``video_reference`` uploads are sent as a base64 data URL: the server + persists them to disk and keeps ``source_path``, handing the pipeline a + real file path instead of a decoded frame list. + """ + return "MiniMax-H3" in (model or "") + + @dataclass class RequestFuncInput: prompt: str @@ -355,13 +367,21 @@ async def async_request_v1_videos( ) elif input.video_paths and len(input.video_paths) > 0: video_path = input.video_paths[0] - image_file = open(video_path, "rb") - form.add_field( - "input_reference", - image_file, - filename=os.path.basename(video_path), - content_type=_guess_mime_type(video_path), - ) + if _is_minimax_h3_mixed_reference(input.model): + with open(video_path, "rb") as reference_file: + video_b64 = base64.b64encode(reference_file.read()).decode("ascii") + form.add_field( + "video_reference", + json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]), + ) + else: + image_file = open(video_path, "rb") + form.add_field( + "input_reference", + image_file, + filename=os.path.basename(video_path), + content_type=_guess_mime_type(video_path), + ) job_id = None job_status = None diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 8f7a02858..6daac9df4 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -858,7 +858,12 @@ class RandomDataset(BaseDataset): height = int(self.args.height or profile.get("height") or 720) num_frames = int(self.args.num_frames or profile.get("num_frames") or 16) fps = float(self.args.fps or profile.get("fps") or 8) - writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)) + # cv2 typically writes MPEG-4 Part 2 (``mp4v``), which pipelines such + # as MiniMax-H3 Ref2VA reject because they require H.264/H.265 + # reference videos. Write the synthetic clip first, then re-encode + # it to H.264 when ffmpeg is available; otherwise keep ``mp4v``. + raw_path = video_path + ".raw.mp4" + writer = cv2.VideoWriter(raw_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)) if not writer.isOpened(): raise RuntimeError("cv2.VideoWriter failed to open") for frame_idx in range(num_frames): @@ -871,6 +876,10 @@ class RandomDataset(BaseDataset): frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220) writer.write(frame) writer.release() + if _reencode_video_to_h264(raw_path, video_path): + os.unlink(raw_path) + else: + os.replace(raw_path, video_path) return [video_path] except Exception as e: raise RuntimeError( @@ -878,6 +887,40 @@ class RandomDataset(BaseDataset): ) from e +def _reencode_video_to_h264(source: str, destination: str) -> bool: + """Re-encode a video to H.264/MP4 using ffmpeg when available. + + Returns False when ffmpeg is missing or the re-encode fails so callers can + fall back to the original container/codec. + """ + import shutil + import subprocess + + if shutil.which("ffmpeg") is None: + return False + try: + result = subprocess.run( + [ + "ffmpeg", + "-y", + "-loglevel", + "error", + "-i", + source, + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + destination, + ], + capture_output=True, + text=True, + ) + return result.returncode == 0 and os.path.exists(destination) + except Exception: + return False + + def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None: """Compute expected execution time (ms) based on a base per-step-per-frame unit time. diff --git a/docs/.nav.yml b/docs/.nav.yml index 07731935b..ca17d550d 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -68,13 +68,12 @@ nav: - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md - - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md + - CPU Offloading: user_guide/diffusion/cpu_offload.md - LoRA: user_guide/diffusion/lora.md - Custom Pipeline: features/custom_pipeline.md - - Request-Level Batching: user_guide/diffusion/request_batching.md - - Step Execution: user_guide/diffusion/step_execution.md + - Execution Modes: user_guide/diffusion/execution_modes.md - Quantization: - - Overview: user_guide/quantization/overview.md + - Quantization: user_guide/quantization/overview.md - Online Quantization: user_guide/quantization/online.md - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md - FP8 W8A8: user_guide/quantization/fp8.md @@ -102,32 +101,49 @@ nav: - CI: contributing/ci - Design Documents: - design/index.md - - design/architecture_overview.md + - Architecture Documents: + - Architecture Overview: design/architecture_overview.md - Feature Design: - - design/feature/disaggregated_inference.md - - design/feature/ray_based_execution.md - - design/feature/omni_connectors/ - - design/feature/prefix_caching.md - - design/feature/cfg_parallel.md - - design/feature/expert_parallel.md - - design/feature/sequence_parallel.md - - design/feature/tensor_parallel.md - - design/feature/vae_parallel.md - - design/feature/hsdp.md - - design/feature/cache_dit.md - - design/feature/teacache.md - - design/feature/skip_softmax.md - - design/feature/async_chunk.md - - design/feature/omni_async_output_materialization.md - - design/feature/vae_parallel.md - - design/feature/diffusion_step_execution.md - - design/feature/diffusion_request_level_batching.md - - design/feature/diffusion_continuous_batching.md + - Runtime and Stage Execution: + - Disaggregated Inference: design/feature/disaggregated_inference.md + - Async Chunk: design/feature/async_chunk.md + - Async Diffusion Output: design/feature/async_diffusion_output.md + - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md + - Automatic Prefix Caching: design/feature/prefix_caching.md + - Communication: + - OmniConnector Implementations: + - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md + - Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md + - Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md + - Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md + - Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md + - Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md + - Diffusion Acceleration: + - Parallelism: + - CFG-Parallel: design/feature/cfg_parallel.md + - Expert Parallel: design/feature/expert_parallel.md + - Hybrid Sharded Data Parallel (HSDP): design/feature/hsdp.md + - Pipeline Parallel: design/feature/pipeline_parallel.md + - Sequence Parallel: design/feature/sequence_parallel.md + - Tensor Parallel: design/feature/tensor_parallel.md + - VAE Patch Parallelism: design/feature/vae_parallel.md + - Attention Backends: + - Skip-Softmax: design/feature/skip_softmax.md + - Quantization: design/feature/quantization.md + - Cache-DiT: design/feature/cache_dit.md + - TeaCache: design/feature/teacache.md + - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md + - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md + - Infrastructure and Performance: + - Prometheus Metrics: design/metrics.md + - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md - Module Design: - - design/module/ar_module.md - - design/module/dit_module.md - - design/module/entrypoint_module.md - - design/module/async_omni_architecture.md + - Runtime Modules: + - AR Module: design/module/ar_module.md + - DIT Module: design/module/dit_module.md + - Orchestration: + - Entrypoint Module: design/module/entrypoint_module.md + - AsyncOmni Architecture: design/module/async_omni_architecture.md - Docs Guide: contributing/DOCS_GUIDE.md - API Reference: - api/README.md diff --git a/docs/community/governance.md b/docs/community/governance.md index 3b0c62c09..63d38cd20 100644 --- a/docs/community/governance.md +++ b/docs/community/governance.md @@ -40,6 +40,7 @@ Committers have write access and merge rights. They typically have deep expertis - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration +- [@NickCao](https://github.com/NickCao): CI, Tests, Entrypoint, and TTS Model Support - [@princepride](https://github.com/princepride): Diffusion and Omni Support - [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs - [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index bdcce41cc..290c15b94 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -387,8 +387,8 @@ that support request-level batching with `step_execution` disabled. Use this together with `max_num_seqs > 1` for bursty serving traffic. `0` disables admission waiting and preserves the lowest first-request latency. -For diffusion request-level batching tuning, see -[Request-Level Batching](../user_guide/diffusion/request_batching.md). +For diffusion execution and batching tuning, see +[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md). Default: `0.0` diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index d390a603a..5e65d2f48 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -1,6 +1,6 @@ **Scope** -For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/). +For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/). **Test Case Design** diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md index 9bdcd29fe..8ff27ee08 100644 --- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md @@ -44,9 +44,11 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [500, 800], - "mean_audio_ttfp_ms": [2000, 3500], - "mean_audio_rtf": [0.25, 0.35] + "H100": { + "mean_ttft_ms": [500, 800], + "mean_audio_ttfp_ms": [2000, 3500], + "mean_audio_rtf": [0.25, 0.35] + } } } ] @@ -142,10 +144,6 @@ pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \ --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \ --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json - -# Optional baseline assertion (default off) -pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \ - --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json ``` See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …). @@ -202,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par | num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. | | request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. | | max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). | -| baseline | object | No | see above | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string). +| baseline | object | No | see below | Optional thresholds. **Hardware-nested only**: top-level keys must be `[hardware-resource]` pytest markers from `pyproject.toml` (`get_hardware_mark_list()`). Under each hardware key, metric names are free-form; values are scalars or sweep-aligned lists. Example: `{"H100": {"throughput_qps": […], "custom_stage_ms": 1.0}, "A3": {…}}`. Runtime filename aliases (`_RUNTIME_DEVICE_ALIASES`) are separate and do not expand this allowlist. | diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 4c01db413..678797652 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -755,7 +755,8 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2) ### Step Execution -See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md) +See the detailed +[Diffusion Continuous Batching design guide](../../design/feature/diffusion_continuous_batching.md). Use this only when your pipeline can be split into stable request-scoped and step-scoped phases. The reference implementation is @@ -769,9 +770,8 @@ step-scoped phases. The reference implementation is Do not enable `step_execution=True` until those four methods are implemented and validated against the request-level path. -If you want the pipeline to work with the experimental batched step-wise path -(`max_num_seqs > 1`), also see: -[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md). +The same design guide covers the experimental batched step-wise path used when +`max_num_seqs > 1`. If you expose this in example scripts or recipes, keep it opt-in. Surface runtime features like `step_execution` as optional flags instead of silently @@ -824,7 +824,7 @@ omni = Omni(model="your-model", ### CPU Offload -See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md) +See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload.md) vLLM-Omni provides two offloading strategies to reduce GPU memory usage: diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md index 482bd1122..2ec89b9d1 100644 --- a/docs/design/feature/diffusion_continuous_batching.md +++ b/docs/design/feature/diffusion_continuous_batching.md @@ -1,165 +1,232 @@ -# Continuous Batching for Step-Wise Diffusion +# Diffusion Continuous Batching -!!! warning "Experimental Feature" - This feature is experimental. It currently applies only to native - diffusion pipelines running with `step_execution=True`. - -This document describes the batching extension built on top of -[Diffusion Step Execution](diffusion_step_execution.md). The base -step-execution contract is unchanged. The batching work is mainly in the -scheduler and runner layers. - -## Why It Helps - -Step-wise execution breaks a long denoise loop into scheduler-visible units. -That gives the runtime a place to admit other compatible requests between -steps instead of waiting for an entire request to finish. - -This matters most in low-MFU or bursty serving scenarios: - -- one request's denoise step may not fully saturate the GPU -- several compatible requests can share the same denoise forward pass -- throughput and device utilization can improve without changing request-local - scheduler state - -This is **not** a guaranteed single-request latency win. The main benefit is -usually higher utilization and better throughput when the workload contains -multiple in-flight compatible requests. +This document describes the unified diffusion execution architecture, including +request execution, request-level batching, step execution, continuous batching, +and output delivery. For user-facing configuration and CLI examples, see +[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md). ## Overview -With continuous batching enabled: +`DiffusionEngine` selects an execution mode from `step_execution` and uses +`max_num_seqs` as that mode's scheduler capacity: -- the scheduler may keep multiple compatible requests active at the same time -- the runner packs request-local step state into one `InputBatch` -- `denoise_step()` runs on that batch -- `step_scheduler()` and `post_decode()` still run per request +| Configuration | Engine mode | Scheduler | Execution | +|---|---|---|---| +| `step_execution=False`, `max_num_seqs=1` | `REQUEST_BATCH` | `RequestScheduler` | One complete request-level `forward()` | +| `step_execution=False`, `max_num_seqs>1` | `REQUEST_BATCH` | `RequestScheduler` | One fused `forward()` over compatible requests | +| `step_execution=True`, `max_num_seqs=1` | `STEP_BATCH` | `StepScheduler` | One request advanced one denoise step per scheduler tick | +| `step_execution=True`, `max_num_seqs>1` | `STEP_BATCH` | `StepScheduler` | Compatible requests advanced together in step waves | -The current implementation is conservative: +Serial request execution is handled inside `REQUEST_BATCH`; it is not a +separate engine mode. Similarly, continuous batching is the multi-request +configuration of `STEP_BATCH`, not an independent execution mode. -- only compatible requests are batched together -- per-request progress and completion remain independent +The engine performs the following initialization: -Here, "continuous batching" means the step-wise path enabled by -`step_execution=True`. Request-mode `DiffusionRequestBatch` is static -request-level batching for one full pipeline `forward()` call; it does not -admit or remove requests between denoise steps. +1. Resolve model-specific pre- and post-processing hooks. +2. Select `REQUEST_BATCH` or `STEP_BATCH`. +3. Construct the configured `DiffusionExecutor`. +4. Construct `RequestScheduler` or `StepScheduler`, unless a scheduler was + explicitly injected. +5. Bind `execute_batch` or `execute_step`. -## Enablement +`DiffusionEngine.make_engine()` owns startup warmup after engine construction. +This keeps construction, backend selection, and warmup failure cleanup in one +lifecycle. -Use `--step-execution` as the feature gate, then increase `--max-num-seqs` -above `1` if you want batching: +## Unified Output Stream -```bash -vllm serve Qwen/Qwen-Image --omni \ - --port 8091 \ - --step-execution \ - --max-num-seqs 8 +All asynchronous diffusion requests use the same per-request output stream, +regardless of whether the caller exposes streaming results to the user. + +```python +async for output in engine.step_streaming(request): + ... ``` -`--max-num-seqs 1` keeps the step-wise path without enabling batching. +When a request is admitted, the engine creates an +`asyncio.Queue[DiffusionOutput]` keyed by its request ID. The busy loop sends +both intermediate chunks and final results through that queue: -For a reproducible replay flow using the bundled serving benchmark, see the -Qwen-Image replay commands in -[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md) -and -[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md). +```text +request + -> scheduler.add_request() + -> scheduler.schedule() + -> execute_batch() or execute_step() + -> scheduler.update_from_output() + -> per-request output queue + -> step_streaming() +``` -## Scheduler +Streaming callers forward every yielded output. Non-streaming callers consume +the same generator to completion and return only its final output. This avoids +maintaining separate future-based and queue-based result paths. -The scheduler derives its batch capacity from `max_num_seqs` through -`max_num_running_reqs`. +`DiffusionEngine.step()` and +`async_add_req_and_wait_for_response()` remain deprecated compatibility +wrappers. New integrations should consume `step_streaming()` or +`async_add_req_and_stream_response()`. -Batch admission is gated by -[`StepBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py), -which is built from shape-sensitive and CFG-sensitive sampling fields. This is -the core correctness rule for batching: requests are only co-batched when they -share the same denoise tensor contract. +### Completion and Cancellation -There are three important details: +Output delivery and scheduler lifecycle are related but separate: -- `num_inference_steps` is not part of the key, so requests with different - total step counts can still share a batch -- requests also do not need to be at the same current denoise progress; active - requests can continue batching even when their current step indices diverge -- admission is still FIFO, so an incompatible request at the head of the - waiting queue blocks later compatible requests +- A terminal scheduler result always finalizes request state. +- A disconnected or cancelled consumer removes only its delivery queue. +- The scheduler may still finish and clean up the request after its queue has + been removed. +- Engine shutdown sends an error output to all remaining streams. -Today that compatibility rule is still shape-sensitive. `height`, `width`, -`num_frames`, and CFG-related fields remain part of the key, so different -resolutions or incompatible guidance settings do **not** co-batch yet. The -key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests -targeting different adapters or scales run in separate batches and the -worker can activate exactly one adapter per step. +This separation prevents cancelled consumers from leaving active request IDs in +the scheduler. -The scheduler batching unit is one logical `OmniDiffusionRequest`. In the -step-wise path, runtime tensor batching is represented as `StepInputBatch`. For -request-mode prompt semantics, see -[Request-Level Batching](../../user_guide/diffusion/request_batching.md). +## Request-Batch Execution -## Runner +Request-batch execution runs the complete diffusion pipeline for each scheduler +wave. Each `OmniDiffusionRequest` remains an independent logical request with +its own prompt, sampling parameters, seed, request ID, output, error, and abort +state. -The runner keeps persistent per-request execution state in -[`StepRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py), -while the scheduler owns a separate lightweight -[`SchedulerRequestState`](gh-file:vllm_omni/diffusion/sched/interface.py) -for queueing and lifecycle tracking. +With `max_num_seqs=1`, `execute_batch()` handles the serial request path. With a +larger value, a pipeline can opt into fused execution by declaring: -For each step, the runner builds an -[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the -active request states: +```python +supports_request_batch = True +``` -- prompt embeddings and masks are normalized and padded -- dynamic tensors such as `latents` and `timesteps` are gathered each step -- buffers are reused when batch composition stays the same +Its `forward()` method must accept `DiffusionRequestBatch` and return one +`DiffusionOutput` per request. The runner validates the result count and maps +outputs back to their original request IDs with `BatchRunnerOutput`. -The step-wise batched path is: +Current in-tree request-batch implementations include the Flux, LTX-2, SD3, and +Qwen-Image pipelines. Treat the pipeline capability flag in the source as +authoritative because support continues to expand. -1. Run `prepare_encode()` for newly admitted requests. +### Request-Batch Data Flow + +- `DiffusionSchedulerOutput` carries scheduled request IDs and payloads. +- `DiffusionRequestBatch` provides the pipeline-facing static batch. +- `BatchRunnerOutput` carries per-request results back to the engine. + +Request-local setup, such as seed handling and output/error mapping, remains +per request. Cache refresh, homogeneous LoRA activation, and +`pipeline.forward(req_batch)` can be shared by the batch. + +### Compatibility and Admission + +`RequestBatchSamplingParamsKey` contains fields that affect the tensor contract, +including shape, guidance, output count, and LoRA identity. Only compatible +requests can share a fused forward call. + +Admission is FIFO and conservative. An incompatible request at the front of the +waiting queue can prevent later compatible requests from joining the current +wave. + +`request_batch_max_wait_ms` optionally delays the first schedule operation of a +new wave so bursty arrivals can accumulate. It applies only when: + +- the pipeline supports request batching; +- step execution is disabled; +- no requests are currently running; and +- the configured wait is greater than zero. + +The wait ends when the queue reaches capacity, stabilizes briefly, reaches its +deadline, or the engine stops. + +## Step-Batch Execution + +Step execution exposes denoising progress to the scheduler. A supporting +pipeline implements four stateful operations: + +| Operation | Responsibility | +|---|---| +| `prepare_encode(state)` | Validate input, encode prompts, initialize latents and timesteps, and create request-local scheduler state | +| `denoise_step(input_batch, *, states=...)` | Run one denoise forward for the scheduler-provided request states | +| `step_scheduler(state, noise_pred)` | Update latents and advance request progress | +| `post_decode(state)` | Decode and postprocess a completed request or output boundary | + +Persistent request state lives in `StepRequestState`. Pipeline-specific fields +that do not belong in the shared contract should be stored in `state.extra`. +Queueing and lifecycle metadata remain in the scheduler's request state. + +Current native pipelines that explicitly enable step execution include +Qwen-Image, HunyuanImage3, and Helios. Step execution alone does not imply +continuous-batching support: Qwen-Image accepts batched step states. +HunyuanImage3 accepts batched step states only when its resolved self-attention +backend is `TORCH_SDPA`; otherwise it rejects groups larger than one request. +Configure `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or +`diffusion_attention_config.default.backend=TORCH_SDPA` when +`max_num_seqs>1`. See the +[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md) +for its validated configuration. Helios supports only a single active step +request and must use `max_num_seqs=1`. + +### Continuous Batching + +When `max_num_seqs>1`, `StepScheduler` can keep multiple compatible requests +active. The runner gathers their state into `InputBatch`, performs one batched +denoise forward, then applies scheduler updates per request: + +1. Receive transferred KV payloads, then run `prepare_encode()` for newly + admitted requests. 2. Build or refresh `InputBatch`. -3. Run one batched `denoise_step(input_batch)`. -4. Slice the batched `noise_pred` back per request. -5. Run per-request `step_scheduler()`. -6. Run `post_decode()` only for requests that finished denoising. -7. Scatter updated latents back into persistent request state with - [`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py). +3. Run one batched `denoise_step(input_batch, states=states)`. +4. Slice noise predictions back to each request. +5. Run each request's `step_scheduler()`. +6. Run `post_decode()` at a chunk boundary or request completion. +7. Scatter updated latents back into persistent request state. -This keeps the shared work limited to the denoise forward pass while preserving -request-local scheduler state and outputs. +`StepBatchSamplingParamsKey` protects the batched tensor contract. Requests can +have different total step counts and current step indices, but shape-sensitive, +CFG-sensitive, and LoRA fields must be compatible. FIFO head-of-line blocking +also applies here. -## Engine +Chunk-capable pipelines may emit intermediate results. Final-only step +pipelines emit when the request finishes even if `streaming_output=True`. -[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides -the background loop and async add-request path needed for multiple requests to -accumulate in the scheduler. +## Model Author Guidelines -When `step_execution=True`, the engine routes work through the step-wise -executor path. The continuous batching behavior is defined by scheduler-side -compatibility gating and runner-side `StepInputBatch` packing. +When converting a request-level pipeline: -## Current Limitations +- Reuse the same helpers as `forward()` to avoid behavior drift. +- Copy request-scoped scheduler state into `state.scheduler`. +- Advance `state.step_index` only in `step_scheduler()`. +- Keep model forward work in `denoise_step()`. +- Keep latent mutation in `step_scheduler()`. +- Keep final decoding equivalent to the tail of `forward()`. +- Store masks, condition latents, and other request-local tensors in the + request state rather than on the shared pipeline. -- Experimental feature; use `max_num_seqs=1` for the conservative single-request - step path. -- Only native pipelines that already support `step_execution=True`. -- Only homogeneous batches keyed by `StepBatchSamplingParamsKey` are supported. -- `cache_backend`, KV transfer, and other request-mode extras are not wired - into the batched step-wise path yet. -- Future work can relax the current same-shape restriction with richer - heterogeneous batching policies such as bucketing or padded execution for - different resolutions. +Before setting `supports_step_execution = True`, validate output parity for the +same seed and sampling parameters, scheduler isolation across concurrent +requests, abort cleanup, progress reporting, and all supported CFG paths. -## Related Files +Before setting `supports_request_batch = True`, validate the batch input +contract, one-output-per-request result shape, request identity and error +mapping, seeded concurrency, LoRA compatibility, and tensor IPC. -- Scheduler base: - [`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py) -- Scheduler interface: - [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py) -- Step scheduler: - [`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py) -- Runner: - [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py) -- Input batch: - [`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py) -- Tests: - [`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py) +## Limitations + +- Request batching requires an explicit pipeline capability when + `max_num_seqs>1`. +- Step execution requires the four-method stateful pipeline contract. +- Both batching paths currently require homogeneous compatibility keys. +- FIFO scheduling can reduce batching opportunities. +- `request_batch_max_wait_ms` trades first-request latency for burst + coalescing. +- All diffusion cache backends are currently unsupported in step mode. +- KV transfer is supported for newly admitted step requests; some other + request-mode extras remain unsupported. +- Step continuous batching remains experimental; use `max_num_seqs=1` for the + conservative step path. + +## Related Implementation + +- Engine: [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py) +- Model contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py) +- Scheduler interface: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py) +- Request scheduler: [`vllm_omni/diffusion/sched/request_scheduler.py`](gh-file:vllm_omni/diffusion/sched/request_scheduler.py) +- Step scheduler: [`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py) +- Runner: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py) +- Request batch: [`vllm_omni/diffusion/worker/request_batch.py`](gh-file:vllm_omni/diffusion/worker/request_batch.py) +- Step input batch: [`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py) diff --git a/docs/design/feature/diffusion_request_level_batching.md b/docs/design/feature/diffusion_request_level_batching.md deleted file mode 100644 index 8c1cdbd33..000000000 --- a/docs/design/feature/diffusion_request_level_batching.md +++ /dev/null @@ -1,175 +0,0 @@ -# Request-Level Batching for Diffusion - -This document describes the request-mode batching path for diffusion pipelines. -For end-user enablement and tuning, see -[Request-Level Batching](../../user_guide/diffusion/request_batching.md). - -This is separate from -[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md). -Request-level batching runs one full pipeline `forward()` over a static batch of -compatible requests. Step-wise continuous batching admits work between denoise -steps when `step_execution=True`. - -## Why It Helps - -The request-level design avoids coupling several logical requests to one request -object. This keeps request identity, abort/error handling, and per-request -metadata unambiguous while still allowing one fused pipeline forward pass for -bursty or concurrent traffic. - -## Overview - -With request-level batching enabled: - -- each `OmniDiffusionRequest` contains one `prompt` and one `request_id` -- the scheduler groups compatible waiting requests into one scheduler wave -- `DiffusionRequestBatch` wraps the scheduled requests for pipeline `forward()` -- batch-capable pipelines return `list[DiffusionOutput]`, one output per - request -- `BatchRunnerOutput` maps each result back to its original `request_id` - -Pipelines opt in with `supports_request_batch = True` and a `forward()` method -that accepts `DiffusionRequestBatch` and returns `list[DiffusionOutput]`. -Pipelines that do not opt in keep the existing per-request execution path. - -## Enablement - -Request-level batching is the request-mode path, so `step_execution` must remain -disabled. Increase `max_num_seqs` above `1` to let the scheduler keep multiple -compatible requests active: - -```bash -vllm serve Qwen/Qwen-Image --omni \ - --port 8091 \ - --max-num-seqs 4 -``` - -For bursty online ingress, `request_batch_max_wait_ms` can add a bounded -admission wait before the first `schedule()` of a scheduler wave: - -```bash -vllm serve Qwen/Qwen-Image --omni \ - --port 8091 \ - --max-num-seqs 4 \ - --request-batch-max-wait-ms 20 -``` - -`request_batch_max_wait_ms=0` disables this wait and is the default. - -## Request Contract - -`OmniDiffusionRequest` represents one logical request. It owns one prompt, -sampling parameters, request id, and request-local metadata. Runtime batches are -formed by the scheduler and represented separately from the request payload. - -Runtime batching is represented by: - -- [`DiffusionSchedulerOutput`](gh-file:vllm_omni/diffusion/sched/interface.py) - for scheduled request ids and request payloads -- [`DiffusionRequestBatch`](gh-file:vllm_omni/diffusion/worker/request_batch.py) - for the pipeline-facing request batch -- [`BatchRunnerOutput`](gh-file:vllm_omni/diffusion/worker/utils.py) for - per-request results - -`DiffusionRequestBatch` intentionally exposes compatibility properties such as -`prompts`, `sampling_params`, `request_id`, and `kv_sender_info` so migrated -pipelines can stay close to upstream code while using a batch-aware contract. - -## Scheduler - -The scheduler derives its capacity from `max_num_seqs` through -`max_num_running_reqs`. It exposes waiting/running queue counters so the engine -can decide whether admission wait is useful before scheduling a new wave. - -Batch compatibility is controlled by -[`RequestBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key -contains shape-sensitive and guidance-sensitive fields, including output count -and LoRA identity. Requests with incompatible shapes, CFG settings, output -counts, LoRA adapters, or LoRA scales are kept in separate batches. - -Admission is conservative: - -- the scheduler only batches compatible requests -- FIFO ordering is preserved -- an incompatible request at the head of the waiting queue blocks later - compatible requests - -## Engine - -[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) resolves -request-batch capability during initialization from the configured pipeline -class, including custom pipeline classes. - -The capability check uses the pipeline class attribute -`supports_request_batch = True`. Pipelines that set this attribute must implement -a request-batch-compatible `forward()` contract and return one -`DiffusionOutput` per request; the runner validates that return shape at runtime. - -When the selected pipeline is batch-capable and `step_execution=False`, request -mode routes scheduler waves through the batch executor path. Otherwise it keeps -the per-request executor path. - -The optional admission wait runs only when: - -- request batching is supported -- `step_execution=False` -- `request_batch_max_wait_ms > 0` -- no requests are currently running - -The wait exits early when the waiting queue reaches capacity, when the queue is -stable for a short window, when the deadline expires, or when the engine stops. - -## Executor And Runner - -The executor exposes two request-mode entries: - -- `execute_request`: one worker call per scheduled request -- `execute_batch`: one worker call for the whole `DiffusionSchedulerOutput` - -On the batch path, the worker builds a `DiffusionRequestBatch` and runs the -pipeline once. Request-local setup remains per request: - -- KV transfer metadata -- random generator and seed handling -- request output/error/abort mapping - -Shared batch setup happens once per batch when possible: - -- cache refresh -- LoRA activation for the homogeneous adapter key -- pipeline `forward(req_batch)` - -Large tensor IPC still uses the shared-memory packing path. The packer traverses -both normal `RunnerOutput.result` wrappers and nested batch results so batched -outputs do not fall back to pickle IPC for tensor payloads. - -## Current Limitations - -- Only pipelines that declare the request-batch contract use fused batch - execution. -- Batches are homogeneous under `RequestBatchSamplingParamsKey`; heterogeneous resolution or - incompatible guidance settings do not co-batch yet. -- FIFO scheduling can reduce batching opportunities when an incompatible - request is at the front of the queue. -- `request_batch_max_wait_ms` improves burst coalescing but can add latency to - the first request in a scheduler wave. Keep it small for latency-sensitive - serving. -- Step-wise continuous batching is documented separately and only applies when - `step_execution=True`. - -## Related Files - -- Request object and request batch: - [`vllm_omni/diffusion/request.py`](gh-file:vllm_omni/diffusion/request.py) -- Scheduler interface: - [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py) -- Scheduler base: - [`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py) -- Engine: - [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py) -- Worker runner: - [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py) -- Executor interface: - [`vllm_omni/diffusion/executor/abstract.py`](gh-file:vllm_omni/diffusion/executor/abstract.py) -- Tests: - [`tests/diffusion/test_diffusion_engine.py`](gh-file:tests/diffusion/test_diffusion_engine.py) diff --git a/docs/design/feature/diffusion_step_execution.md b/docs/design/feature/diffusion_step_execution.md deleted file mode 100644 index 88fddb215..000000000 --- a/docs/design/feature/diffusion_step_execution.md +++ /dev/null @@ -1,131 +0,0 @@ -# Adding Step Execution Support for Diffusion Pipelines - -This guide documents vLLM-Omni's stepwise diffusion contract for model authors -and contributors implementing `step_execution=True` support for a diffusion -pipeline. - -For end-user enablement, supported models, and current limitations, see -[Step Execution](../../user_guide/diffusion/step_execution.md). - -This document describes the base step-execution contract only. For the -experimental batching policy layered on top of the step-wise path, see -[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md). - -## Current Support Scope - -`step_execution` is **not** a generic diffusion toggle. It only works for -pipelines that implement the segmented stateful contract in -[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py). - -This page is intentionally author-facing. Treat runtime enablement -(`step_execution=True` in Python or `--step-execution` in serving) as an -opt-in user knob layered on top of the implementation contract below. - -Current in-tree support: - -| Pipeline | Example models | Step execution | -|----------|----------------|----------------| -| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | -| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No | - -Current engine/runtime limitations: - -- Continuous batching with `max_num_seqs > 1` is experimental and documented in - [Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md). - Keep `max_num_seqs=1` for the conservative single-request step path. -- `cache_backend` is not supported in step mode. -- Request-mode extras such as KV transfer are not wired into step mode yet. -- Unsupported pipelines now fail early during model loading instead of failing on the first request. - -## Execution Contract - -Step mode is driven by four pipeline methods plus the shared mutable request -state object: - -- `prepare_encode(state)`: one-time request preparation. -- `denoise_step(state)`: compute the noise prediction for the current step. -- `step_scheduler(state, noise_pred)`: mutate latents and advance step state. -- `post_decode(state)`: decode the final output after denoising is complete. - -The state lives in -[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py) -as `StepRequestState`. Store request-scoped tensors there, or use -`state.extra` for model-specific fields that do not justify extending the core -dataclass. - -The worker-side step loop lives in -[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py): - -1. `prepare_encode()` runs once for a new request. -2. `denoise_step()` runs every scheduler tick. -3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`. -4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true. - -## Recommended Split - -When converting an existing request-level `forward()` pipeline, keep the split -strict and mechanical: - -| Request-level phase | Stepwise method | What belongs there | -|---------------------|-----------------|--------------------| -| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request | -| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep | -| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step | -| VAE decode / postprocess | `post_decode()` | Final decode only | - -Keep the stepwise path reusing the same helpers as the request-level path -whenever possible. Reimplementing the denoise loop from scratch is the easiest -way to introduce behavioral drift. - -## Qwen-Image Reference - -[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py) -is the reference implementation and is split correctly for the current -contract: - -- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding, - latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned - with `forward()`. -- `prepare_encode()` deep-copies `self.scheduler` **after** - `prepare_timesteps()` so request-specific scheduler state is isolated. -- `denoise_step()` reuses `_build_denoise_kwargs()` plus - `predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and - non-CFG behavior stay identical to the request-level path. -- `step_scheduler()` only calls - `scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)` - and increments `state.step_index`. -- `post_decode()` reuses `_decode_latents()`, so the final image decode matches - the normal `forward()` path. - -That decomposition is the target pattern for future models. - -## Rules For New Pipelines - -- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into - `state.scheduler` during `prepare_encode()`. -- Do not mutate `state.step_index` inside `denoise_step()`. Only - `step_scheduler()` should advance the step. -- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`. -- If the request-level pipeline has condition latents, masks, or edit-specific - tensors, store them in `state` or `state.extra`, not in global pipeline - attributes. -- Preserve CFG behavior by sharing the same helper path used by `forward()`. -- Keep `post_decode()` equivalent to the tail of `forward()`. - -## Validation Checklist - -Before marking a pipeline as `supports_step_execution = True`, verify: - -- Stepwise output matches request-level output for the same seed and sampling params. -- Per-request scheduler state is isolated across concurrent requests. -- Abort during denoise does not leak cached state. -- `step_index` reported by `RunnerOutput` matches the scheduler progress. -- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them. - -## Related Files - -- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py) -- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py) -- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py) -- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py) -- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py) diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md new file mode 100644 index 000000000..da735d60b --- /dev/null +++ b/docs/design/feature/distributed_layerwise_offload.md @@ -0,0 +1,156 @@ +# Distributed Layerwise Offload + +This document describes distributed layerwise offload (DLO) for diffusion +models. DLO keeps only a small number of DiT blocks on the accelerator and +streams the remaining blocks from host memory. The distributed backend can +either shard those host-side weights across an existing parallel group or keep +the standard loader's rank-local weights and avoid an additional collective. + +For user-facing commands, see the +[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md) +and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md). + +## Status + +DLO is implemented for multi-device diffusion execution. The default +AllGather path is the primary path for DP and SP deployments. The +`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful +for standard-loader sharding, workstation bring-up, and systems where an +additional DLO collective is undesirable, but it does not reduce host weight +storage across ranks. + +The compatibility matrix below describes the current implementation. The +unit-level guards are covered, but not every parallelism combination has a +full model-and-hardware end-to-end test. + +## Design + +### DLO consumes the existing parallel topology + +DLO does not create a new DP, TP, or SP topology. It reads the configured +`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the +standard distributed groups have been initialized. + +The DLO weight-sharding group is selected as follows: + +1. Use the existing DP group when `data_parallel_size > 1`. +2. When DP is one and SP is greater than one, use the SP group. +3. Otherwise, run rank-locally without a DLO process group. + +TP is deliberately not used as DLO's AllGather group. HSDP has its own +parameter-sharding lifecycle and is not allowed to be sharded a second time by +DLO's AllGather path. + +### AllGather path + +With the default `dlo_use_allgather=True`, each rank stores approximately +`1 / group_size` of each streamable block in pinned host memory. The next +block's shard is copied to a device buffer and reconstructed with +`all_gather_into_tensor` on a communication stream while the current block is +executing. + +```text +Compute: [Block N] [Block N+1] [Block N+2] +H2D: [shard N+1] [shard N+2] +AllGather: [full N+1] [full N+2] +Buffers: [current slot] [prefetch slot] [current slot] +``` + +![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif) + +The backend uses two shared device buffers, so accelerator weight residency is +bounded by the largest streamed blocks rather than the complete model. + +When DP is greater than one, the engine can process one request per DP rank in +the same denoising wave. Because AllGather is a collective, all participating +requests must take the same execution path at every denoising step. + +### Rank-local path without DLO AllGather + +With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to +one. The regular model loader remains responsible for preparing each rank's +weights, including TP-local tensors or HSDP-managed parameters. DLO then +streams those rank-local tensors block by block using H2D copies only. + +This mode means: + +- DP still provides independent replicas, but DLO does not shard weights + across DP ranks. +- SP still performs its normal activation/attention collectives, but DLO does + not shard weights across SP ranks. +- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only + DLO's additional weight AllGather is disabled. +- Pure DP deployments keep a full host-side model copy per rank, subject to + shared OS page-cache behavior. +- The scheduler does not require a synchronized DP request wave for DLO. + +This path is intentionally not implemented by reusing the AllGather mmap +loader. It relies on the standard loader so model-specific TP/HSDP transforms +remain intact. + +## Parallelism compatibility + +| Parallelism | DLO + AllGather | DLO without AllGather | +|---|---|---| +| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. | +| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. | +| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. | +| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. | + +### Combined dimensions + +- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater + than one; SP continues to use its own sequence-parallel group. If DP is one, + the SP group becomes DLO's sharding group in AllGather mode. +- **DP + TP/SP without AllGather:** standard model loading defines the + rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight + collective. +- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but + DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight + materialization and synchronization. +- **HSDP + DP or TP:** rejected independently by the diffusion parallel + configuration. + +## Request and loading constraints + +AllGather DP multi-concurrency requires: + +- explicit `num_inference_steps`; +- the same `num_inference_steps` for all requests in a wave; and +- identical request arguments that affect the collective execution path. + +The no-AllGather path does not impose these DLO-specific synchronized-wave +requirements. + +The mmap loader is used only by the supported DLO+AllGather path when the +model has a compatible checkpoint layout. Online quantization is incompatible +with that sharded mmap path; use `--dlo-no-use-allgather` or disable online +quantization. Models that do not meet the mmap requirements use the regular +loader path. + +## Validation coverage + +Current source-level validation includes: + +- HSDP + DLO + AllGather rejection; +- HSDP + DLO without AllGather acceptance at configuration level; +- TP rejection in the DLO+AllGather mmap path; +- resident-layer requests requiring no-AllGather; +- DP request-wave validation for denoising-step compatibility; +- sharding, double-buffer, AllGather-size, and heterogeneous-block regression + tests. + +The highest-value missing coverage is end-to-end numerical comparison against +ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather +on the target CUDA/NCCL or CANN/HCCL hardware. + +## Recommendations + +- Use **DP + DLO AllGather** for the supported throughput and host-memory + scaling path. +- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is + not the goal. +- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with + the expectation of higher host-memory use and lower validation confidence. +- Prefer **HSDP alone** for production HSDP deployments until the combined + HSDP + DLO no-AllGather path has broader end-to-end coverage. diff --git a/docs/design/feature/quantization.md b/docs/design/feature/quantization.md new file mode 100644 index 000000000..37dd5e26b --- /dev/null +++ b/docs/design/feature/quantization.md @@ -0,0 +1,249 @@ +# Quantization + +This document describes the quantization architecture in vLLM-Omni and the +extension points for contributors. For user-facing configuration, supported +hardware, and method-specific instructions, see the +[quantization user guide](../../user_guide/quantization/overview.md). + +## Goals and scope + +The design has four goals: + +1. Provide one `quantization_config` entry point for diffusion-only, + multi-stage omni/TTS, and multi-stage diffusion models. +2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where + a vLLM backend already provides the required behavior. +3. Allow different components of one pipeline to use different quantization + configurations. +4. Preserve the distinction between online quantization of a BF16/FP16 + checkpoint and loading a pre-quantized checkpoint with serialized scales. + +This document covers configuration resolution, component routing, model +integration, backend extension, and validation. Kernel implementations and +user-facing support matrices are outside its scope. + +## Architecture + +The quantization path is split across the following layers: + +| Layer | Responsibility | Source | +| --- | --- | --- | +| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` | +| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` | +| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` | +| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` | +| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` | +| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` | + +The runtime contract remains vLLM's `QuantizationConfig`: + +```text +user/CLI or model config + | + v +OmniDiffusionConfig.quantization_config + | + +--> build_quant_config() + | | + | +--> Omni override + | +--> vLLM registry backend + | +--> ComponentQuantizationConfig + | + +--> checkpoint metadata reconciliation + | + +--> diffusion worker / vLLM config + | + +--> model layers call get_quant_method(layer, prefix) +``` + +The model layer receives either a method implementation or `None`. Returning +`None` leaves that layer unquantized. + +## Configuration resolution + +`build_quant_config()` accepts the following forms: + +| Input | Resolution | +| --- | --- | +| `None` or `"none"` | Disable quantization. | +| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. | +| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. | +| Per-component dictionary | Construct a `ComponentQuantizationConfig`. | +| Existing `QuantizationConfig` | Pass through without rebuilding it. | + +Method aliases are normalized case-insensitively, with `-` and `_` treated as +equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter. +ModelOpt checkpoint metadata is detected from its method, producer metadata, or +`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration. + +The factory uses Omni-specific overrides for methods whose diffusion behavior +or platform dispatch is not provided by the vLLM registry. The current override +set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/ +AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through +vLLM's registry when their configuration is compatible. + +## Per-component routing + +Multi-stage models do not necessarily want one quantization policy for every +module. A component map expresses that policy without requiring each pipeline +to construct separate vLLM configurations: + +```python +quantization_config = { + "transformer": {"method": "fp8"}, + "vae": None, + "default": None, +} +``` + +`ComponentQuantizationConfig` applies the following rules: + +1. Component keys are interpreted as runtime layer-prefixes. +2. The longest matching prefix wins. +3. `None` disables quantization for the matching component. +4. `default`, when present, handles prefixes with no explicit match. +5. A matching child config delegates to its normal + `get_quant_method(layer, prefix)` implementation. + +Prefix matching happens after the model's weight-name mapping. New model +integrations must verify the actual runtime prefixes, especially when a +`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route +a layer to the default configuration, so component-routing tests should cover +both positive matches and fall-through behavior. + +## Online and pre-quantized checkpoints + +The same method name can represent two different loading paths: + +- **Online**: the source checkpoint contains BF16/FP16 weights; the method + creates quantized weights or scales while the model is loaded. +- **Pre-quantized**: the checkpoint contains quantized weights and the metadata + required by the method; the loader consumes those serialized tensors. + +Method configuration owns the mode flag, for example +`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The +configuration's `get_quant_method()` selects the corresponding linear method. + +For pipelines with transformer-specific `config.json` files, such as cascade +models, `resolve_quant_config_from_disk()` reconciles the active configuration +with each transformer's checkpoint metadata. Its invariants are: + +- If no active configuration exists, valid checkpoint metadata can be used for + auto-detection. +- A method mismatch is an error rather than an implicit conversion. +- Serialized-checkpoint flags rebuild an online configuration into the + matching offline configuration. +- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active + configuration when they differ. +- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path. + +This prevents a model from interpreting serialized weights with an online +loader or silently applying the wrong method to one transformer in a +multi-transformer pipeline. + +## Model and layer integration + +Diffusion model constructors receive the resolved configuration and pass it to +quantizable vLLM layers. Each method configuration is responsible for: + +- identifying supported layers in `get_quant_method()`; +- creating the correct parameter shapes and loaders; +- processing weights after loading when conversion is required; and +- applying the platform-specific quantized operation at runtime. + +Model code must keep precision-sensitive or unsupported modules unquantized. +The shared `safe_quant_config()` helper is used by model integrations for +normalization and modulation layers where a generic FP8 configuration would be +unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers +are not implicitly quantized just because the main transformer is quantized; +they require explicit component routing and method support. + +Weight quantization and quantized KV cache are separate paths. The former is +resolved through `quantization_config`; runtime attention/KV-cache quantization +is configured by the diffusion attention configuration and must be validated +independently. + +## Platform and parallelism boundaries + +Platform checks belong in the method implementation rather than in the common +factory. This keeps configuration parsing portable while allowing CUDA, NPU, +and XPU methods to select different kernels or reject unsupported execution at +model construction time. The user guide is the source of truth for the current +hardware support matrix. + +Quantization methods must also respect the weight layout produced by the +configured parallelism strategy: + +- Tensor-parallel layers may receive partitioned weight shapes and loaders. +- Sequence-parallel execution changes activation movement but does not by + itself change the quantization configuration. +- Data-parallel replicas build their own method state; quantization metadata + must be available consistently on every rank. +- `WeightsMapper` transformations must be applied before matching ignored-layer + lists or component prefixes. + +Distributed layerwise offload has an additional boundary: its sharded mmap +AllGather path is not compatible with online quantization because the path +expects already materialized checkpoint weights. Use the rank-local +`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when +combining these features. See the [DLO design document](distributed_layerwise_offload.md) +for the full compatibility matrix. + +## Adding a quantization backend + +Contributors should follow this sequence: + +1. Check whether an existing vLLM backend can provide the required weight + format and kernel behavior. Prefer reuse over a duplicate Omni backend. +2. If Omni-specific behavior is required, implement a `QuantizationConfig` and + its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`. +3. Implement `get_name()`, supported activation dtypes, minimum capability, + checkpoint parsing, ignored-layer handling, and `get_quant_method()`. +4. Add a lazy factory override in `factory.py`. Keep optional platform + dependencies out of package import paths and normalize aliases to one + canonical method name. +5. If the method is pre-quantized, define the checkpoint metadata and the + serialized loading path before adding an online path. +6. Add model-specific component routing only when the model contains stages + that must use different policies. +7. Document user-visible configuration in the quantization user guide and keep + implementation rationale here. + +The public factory registration hook, +`register_quantization_override(method, builder)`, is intended for integrations +that need to register an Omni-compatible builder without changing the vLLM +registry. + +## Validation + +Validation should cover the complete path, not only whether a config object can +be instantiated: + +| Area | Required checks | +| --- | --- | +| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps | +| Method config | Online/offline selection, ignored layers, platform guards, and weight loading | +| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata | +| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts | +| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts | +| Performance | Peak device memory, generation latency, and any backend-specific throughput metric | + +The unit tests live under `tests/diffusion/quantization/`. The trajectory +similarity tool in +`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` +provides a repeatable quality and memory comparison for diffusion models. +Method support and feature combinations remain model- and hardware-dependent; +passing factory tests alone does not establish production support. + +## Design invariants + +- `quantization_config` is the single runtime source of truth after config + normalization. +- A checkpoint method mismatch fails loudly instead of risking weight + corruption. +- Components without an explicit supported quantization path remain in their + original precision. +- Optional backend dependencies are imported lazily. +- Method implementations own platform-specific capability checks and kernels. +- User-facing support claims require model, hardware, numerical, and memory + validation in addition to configuration-level tests. diff --git a/docs/design/feature/ray_based_execution.md b/docs/design/feature/ray_based_execution.md deleted file mode 100644 index 4f4caa642..000000000 --- a/docs/design/feature/ray_based_execution.md +++ /dev/null @@ -1,59 +0,0 @@ -# Distributed utils - -This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends. -## 1. Installation -```bash -pip install "ray[default]" -``` -## 2. Ray Utils - -The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for: -* **Multi-node deployment**: Running pipeline stages across different physical machines. -* **Resource management**: Efficient GPU/CPU allocation. - -### 2.1 Basic Usage - -To use the Ray backend, specify `worker_backend="ray"` when initializing the engine. - -**Command Line Example:** -```bash -vllm serve Qwen/Qwen2.5-Omni-7B \ - --omni \ - --port 8091 \ - --worker-backend ray \ - --ray-address auto -``` - -### 2.2 Cluster Setup - -**Step 1: Start Head Node** -Run this on your primary machine: -```bash -ray start --head --port=6399 -``` - -**Step 2: Connect Worker Nodes** -Run this on each worker machine: -```bash -ray start --address=:6399 -``` - -> **Tip**: For a complete cluster setup script, refer to the vLLM example: -> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh) - -### 2.3 Distributed Connector Support - -When running on Ray, the system automatically adapts its communication strategy: - -* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback). -* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma). -* **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`. - -### 2.4 Internal Helpers - -* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one. - -## 3. Troubleshooting - -* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open. -* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python. diff --git a/docs/design/figures/dlo/dlo_pipeline.gif b/docs/design/figures/dlo/dlo_pipeline.gif new file mode 100644 index 000000000..cc206d03b Binary files /dev/null and b/docs/design/figures/dlo/dlo_pipeline.gif differ diff --git a/docs/design/index.md b/docs/design/index.md index 593191616..5434b27fb 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -1,6 +1,8 @@ # Design Documents -This section contains design documents and architecture specifications for vLLM-Omni. +This section contains design documents and architecture specifications for +vLLM-Omni. The sidebar groups documents by the system concern they describe, +while the source paths remain stable so existing links continue to work. ## Architecture Documents @@ -8,19 +10,63 @@ This section contains design documents and architecture specifications for vLLM- ## Feature Design Documents -- [Disaggregated Inference](feature/disaggregated_inference.md) -- [Ray-based Execution](feature/ray_based_execution.md) -- [Async Omni Output Materialization](feature/omni_async_output_materialization.md) -- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md) -- [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md) -- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md) +### Runtime and stage execution -## Infrastructure Design Documents +- [Disaggregated Inference](feature/disaggregated_inference.md) +- [Async Chunk](feature/async_chunk.md) +- [Async Diffusion Output](feature/async_diffusion_output.md) +- [Async Omni Output Materialization](feature/omni_async_output_materialization.md) +- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md) + +### Communication + +#### OmniConnector implementations + +- [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md) +- [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md) +- [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md) +- [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md) +- [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md) +- [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md) + +### Diffusion acceleration + +#### Parallelism + +- [CFG-Parallel](feature/cfg_parallel.md) +- [Expert Parallel](feature/expert_parallel.md) +- [Hybrid Sharded Data Parallel (HSDP)](feature/hsdp.md) +- [Pipeline Parallel](feature/pipeline_parallel.md) +- [Sequence Parallel](feature/sequence_parallel.md) +- [Tensor Parallel](feature/tensor_parallel.md) +- [VAE Patch Parallelism](feature/vae_parallel.md) + +#### Attention Backends + +- [Skip-Softmax](feature/skip_softmax.md) + +#### Quantization + +- [Quantization](feature/quantization.md) + +- [Cache-DiT](feature/cache_dit.md) +- [TeaCache](feature/teacache.md) +- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md) +- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md) + +## Infrastructure and Performance - [Prometheus Metrics](metrics.md) +- [Speech Generation Performance Optimizations](qwen3_omni_tts_performance_optimization.md) ## Module Design Documents +### Runtime Modules + - [AR Module](module/ar_module.md) - [DIT Module](module/dit_module.md) + +### Orchestration + - [Entrypoint Module](module/entrypoint_module.md) +- [AsyncOmni Architecture (Qwen3-Omni Example)](module/async_omni_architecture.md) diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index aabdb78b0..36fd16877 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -252,8 +252,7 @@ while True: batch of compatible independent requests for one full pipeline forward when the pipeline supports request-level batching; the step-wise path can also admit or remove compatible requests between denoise steps. See - [Request-Level Batching](../feature/diffusion_request_level_batching.md) and - [Continuous Batching for Step-Wise Diffusion](../feature/diffusion_continuous_batching.md). + [Diffusion Continuous Batching](../feature/diffusion_continuous_batching.md). --- @@ -540,7 +539,7 @@ def get_attn_backend_for_role(role, head_size, attention_config=None, role_categ 4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware. -A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today the only such field is `skip_softmax` (TRTLLM_ATTN), so e.g. `target_sparsity` can be set without changing model code. +A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code. For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md). diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 15ceb395e..f9673eac7 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -87,9 +87,8 @@ if __name__ == "__main__": !!! info - For diffusion request-level batching controls such as `max_num_seqs` and - `request_batch_max_wait_ms`, see - [Request-Level Batching](../user_guide/diffusion/request_batching.md). + For diffusion request batching, step execution, and streaming controls, see + [Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md). For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md) diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 92a13aa74..7dc35b367 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -14,7 +14,7 @@ The full set of backends and their platform defaults is in the **Backend Options | Value | Notes | |---|---| -| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | +| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | | `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. | | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). | | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. | @@ -23,6 +23,7 @@ The full set of backends and their platform defaults is in the **Backend Options | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. | | `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. | | `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. | +| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. | ## Configuration @@ -69,8 +70,9 @@ vllm-omni serve \ --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}' ``` -A backend that needs configuration exposes it as a typed field on the spec. Today the only -one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)): +A backend that needs configuration exposes it as a typed field on the spec: +`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and +`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)): ```bash --diffusion-attention-config.default.backend TRTLLM_ATTN \ @@ -103,7 +105,7 @@ A plain dict is also accepted and normalized to `AttentionConfig`. Auto-route preference, in order: -1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model is mask-free** (see below) +1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path** 2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels) 3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5 4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel @@ -111,7 +113,7 @@ Auto-route preference, in order: `TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below. -`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell because it is the only backend that can enable Skip-Softmax, not because its dense kernel is faster — dense, the two are comparable. It cannot honor an attention mask, so the auto-default only applies to pipelines verified mask-free (the Wan family); mask-using pipelines keep `CUDNN_ATTN`. `TRTLLM_ATTN` stays available everywhere via explicit `--diffusion-attention-backend TRTLLM_ATTN` (which raises clearly if a mask is received). +`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback. The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route. @@ -204,11 +206,80 @@ AttentionConfig( ) ``` - **Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103** (B300); `int8` QK kernels are compiled for **SM100 only**. +## RAINFUSION_ATTN Backend and Block-Sparse Video Attention + +`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each +128-token block of keys, ranks them per query block, and computes attention against only the +top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are +spatiotemporal neighbours rather than arbitrary rows of the packed sequence. + +Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the +first-frame blocks are always kept dense, which is why the realized sparsity is lower than the +nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an +exempt layer, a sequence with no published video geometry, a video segment under 32 blocks — +delegates to `FLASH_ATTN`, so a model can select this backend unconditionally. + +Enable it through the typed `block_sparse` block on the attention spec, shared by every +block-sparse backend: + +| Key | Valid values | Meaning | +|---|---|---| +| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. | +| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. | +| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. | + +```bash +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN", + "block_sparse": {"sparsity": 0.8, "start_step": 0}}}' +``` + +Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction): + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec + +AttentionConfig( + default=AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"), + ), +) +``` + +**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the +cheapest way to recover structure, because the early high-noise steps decide global layout while the +later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for +`skip_layers` last, once an A/B against dense at the same seed points at specific blocks. + +Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also +incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank +blocks — use Ulysses SP (`ring_degree=1`). + +### Which geometries run sparse + +Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count +is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own +tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows +that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense +attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no +speedup, so pick an aligned geometry when you want one. + +Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8 +spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact +patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns +are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially +distant positions and selection can no longer rank them meaningfully. This is silent, and invisible +to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped. + +"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned +resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`, +which you can buy back with `start_step`. + ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell) Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`. diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload.md similarity index 88% rename from docs/user_guide/diffusion/cpu_offload_diffusion.md rename to docs/user_guide/diffusion/cpu_offload.md index 86a0e9201..02ab9f074 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload.md @@ -184,10 +184,15 @@ class Flux2Transformer2DModel(nn.Module): ### How It Works Distributed layerwise offloading extends single-GPU layerwise offloading to -multi-device deployments. Each DP rank stores only **1/dp_size** of the model -weights in host memory; full layer weights are reconstructed at runtime via -**AllGather** on a dedicated communication stream, overlapped with computation -via a fixed double-buffer scheme. +multi-device deployments. In the default mode, each rank stores only +**1/dp_size** of the model weights in host memory; full layer weights are +reconstructed at runtime via **AllGather** on a dedicated communication stream, +overlapped with computation via a fixed double-buffer scheme. The +`--dlo-no-use-allgather` mode instead streams standard-loader rank-local +weights independently and does not shard weights across ranks. + +For the implementation design and compatibility matrix, see +[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md). **Key features:** - **Weight sharding + AllGather**: each rank stores 1/dp_size of weights, @@ -219,7 +224,7 @@ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --data-parallel-size 4 -# Without AllGather (each rank loads full weights, no sharding) +# Without DLO AllGather (each rank streams standard-loader rank-local weights) vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --data-parallel-size 4 \ @@ -247,9 +252,9 @@ m = Omni( | Flag | Description | Default | |------|-------------|---------| | `--enable-distributed-layerwise-offload` | Enable DLO | `false` | -| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` | +| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` | | `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` | -| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` | +| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` | ### How model weights are loaded (mmap path) @@ -303,9 +308,15 @@ request while AllGather synchronizes only weight shards (request-independent). ### Limitations - Online quantization (FP8) is incompatible with mmap loading — falls back to - regular `load_weights()` automatically -- Tensor Parallel is not supported (DLO uses DP-based sharding) -- HSDP + AllGather is rejected (would double-shard weights) + regular loading only when the model is outside the DLO+AllGather mmap path; + online quantization with DLO+AllGather is rejected. Use + `--dlo-no-use-allgather` or disable online quantization. +- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware + weight-loader callbacks are bypassed. The no-AllGather path retains the + standard TP loader and is experimental. +- HSDP + AllGather is rejected because it would double-shard parameters. HSDP + with no-AllGather is accepted at configuration level, but has limited + end-to-end validation. - `num_inference_steps=None` is not allowed in DP multi-concurrency mode **Module Discovery** @@ -338,7 +349,7 @@ OffloadBackend (base class) ├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook │ (single-GPU, full weights on host) └── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook - (multi-GPU, 1/dp_size sharded weights + AllGather) + (multi-GPU, sharded AllGather or rank-local mode) ``` Factory function `get_offload_backend()` selects the appropriate backend based on @@ -369,4 +380,10 @@ reasoner/generator components inside the model forward pass. **Notes:** - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks -- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples. +- Distributed Layerwise Offloading reuses the layerwise block topology, but the + model and parallelism combination still needs validation. AllGather mode + additionally requires a compatible checkpoint/mmap layout; no-AllGather mode + retains the standard loader's rank-local tensors. See the + [DLO design](../../design/feature/distributed_layerwise_offload.md) and + [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) + for usage examples. diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md new file mode 100644 index 000000000..9611a8210 --- /dev/null +++ b/docs/user_guide/diffusion/execution_modes.md @@ -0,0 +1,252 @@ +# Diffusion Execution Modes + +vLLM-Omni supports complete-request execution and step-wise diffusion +execution. Both modes use the same asynchronous engine output stream; the +configuration controls how work is scheduled and whether intermediate outputs +are exposed. + +## Choose a Mode + +| Goal | CLI configuration | +|---|---| +| Serial request execution | `--max-num-seqs 1` | +| Fused request-level batching | `--max-num-seqs N` | +| Single-request step execution | `--step-execution --max-num-seqs 1` | +| Step-wise continuous batching | `--step-execution --max-num-seqs N` | +| Chunked diffusion output | `--diffusion-streaming-output` | + +`N` must be greater than `1` to allow batching. The selected pipeline must +support request-level batching or batched-step execution as appropriate; +single-request step support alone is not sufficient for step-wise continuous +batching. + +## Request Execution + +Request execution is the default when `--step-execution` is omitted. The +pipeline performs a complete `forward()` for each scheduler wave. + +### Serial Requests + +Use the conservative serial path with any request-mode pipeline: + +```bash +vllm serve MODEL --omni \ + --port 8091 \ + --max-num-seqs 1 +``` + +### Request-Level Batching + +Set `max_num_seqs` above one to combine compatible independent requests into a +single pipeline forward: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --max-num-seqs 4 +``` + +For bursty traffic, add a small admission window: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --max-num-seqs 4 \ + --request-batch-max-wait-ms 20 +``` + +`--request-batch-max-wait-ms 0` is the default. A nonzero value can improve +batch formation but adds up to that much latency before a new scheduler wave. +It has no effect in step mode. + +Each prompt remains a separate logical request. Do not submit a top-level list +as one packed prompt. Concurrent serving requests are batched internally when +their shapes, guidance settings, output counts, and LoRA settings are +compatible. + +Only pipelines that declare request-batch support accept +`max_num_seqs>1` in request mode. Unsupported pipelines fail during engine +initialization; use `max_num_seqs=1` instead. + +## Step Execution + +Step execution lets the scheduler advance and abort requests between denoise +steps: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --step-execution \ + --max-num-seqs 1 +``` + +Set a larger capacity to allow compatible requests to share denoise waves: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --step-execution \ + --max-num-seqs 8 +``` + +Step continuous batching is experimental. Start with `max_num_seqs=1` when +validating a model or debugging correctness, then increase it for +multi-request throughput. + +Step execution is capability-based, not a generic switch for every diffusion +model. Qwen-Image supports step-wise continuous batching. HunyuanImage3 also +supports it, but only when its resolved self-attention backend is `TORCH_SDPA`; +set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure +`diffusion_attention_config.default.backend=TORCH_SDPA` before using +`--max-num-seqs >1`. See the +[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md) +for its validated configuration. Helios supports single-request step execution only: use +`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's +documentation and source for the latest support status. + +## Streaming Output + +Use `--diffusion-streaming-output` for a pipeline that can produce intermediate +diffusion outputs: + +```bash +vllm serve BestWishYsh/Helios-Distilled --omni \ + --port 8000 \ + --diffusion-streaming-output +``` + +Streaming output requires step execution. If +`--diffusion-streaming-output` is set without `--step-execution`, the engine +enables step execution automatically. Model initialization fails if the +pipeline does not implement step execution. + +Chunk-capable pipelines emit intermediate and final outputs through the same +request stream. Final-only step pipelines emit only the final result even when +streaming output is enabled. + +The normal non-streaming serving path uses the same internal stream but drains +it and returns only the final output. Users do not need to enable +`--diffusion-streaming-output` to benefit from the unified engine lifecycle. + +## Send Requests + +Execution-mode flags configure the server. They do not change the +OpenAI-compatible client request format. For example, after starting any of +the non-streaming configurations above: + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "a cat sitting on a windowsill", + "size": "1024x1024", + "num_inference_steps": 50, + "seed": 42 + }' | jq -r '.data[0].b64_json' | base64 -d > cat.png +``` + +To create a request-level batch, send multiple independent requests +concurrently. Do not put several prompts into one `prompt` field: + +```bash +for prompt in "a red fox" "a blue bird" "a green frog"; do + curl -s -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d "{\"prompt\":\"${prompt}\",\"size\":\"1024x1024\",\"seed\":42}" \ + > /dev/null & +done +wait +``` + +The scheduler may batch these requests when their sampling parameters are +compatible. See the +[Image Generation API](../../serving/image_generation_api.md) for response +formats and additional client examples. Streaming models can use +model-specific streaming endpoints documented by their serving guide. + +## Python API + +The Python arguments mirror the CLI flags: + +```python +from vllm_omni import Omni + +omni = Omni( + model="Qwen/Qwen-Image", + max_num_seqs=4, + request_batch_max_wait_ms=20.0, +) + +outputs = omni.generate( + [ + "a cup of coffee on a table", + "a toy dinosaur on a sandy beach", + "a fox waking up in bed and yawning", + ] +) +``` + +Each list item passed to `Omni.generate()` becomes an independent logical +request that the scheduler may batch with compatible items. + +For step execution: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni( + model="Qwen/Qwen-Image", + step_execution=True, + max_num_seqs=1, +) + +outputs = omni.generate( + "A cat sitting on a windowsill", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +## Stage Configuration + +Use the equivalent engine arguments in a deployment YAML: + +```yaml +stage_args: + - stage_id: 0 + stage_type: diffusion + engine_args: + step_execution: false + max_num_seqs: 4 + request_batch_max_wait_ms: 20 +``` + +For step execution, set `step_execution: true` and remove +`request_batch_max_wait_ms`. + +## CLI Reference + +| Flag | Default | Effect | +|---|---:|---| +| `--step-execution` | disabled | Select step-wise scheduling | +| `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity | +| `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode | +| `--diffusion-streaming-output` | disabled | Expose supported intermediate diffusion outputs and require step execution | + +## Limitations and Troubleshooting + +- Requests batch only when their compatibility-sensitive parameters match. +- Different LoRA adapters or scales run in separate batches. +- FIFO scheduling can cause an incompatible request to block later compatible + requests. +- All diffusion cache backends are unsupported in step mode. KV transfer is + supported for newly admitted step requests; some other request-mode extras + remain unsupported. +- If request-mode startup reports that the pipeline does not support batching, + use `--max-num-seqs 1`. +- If step-mode startup mentions `prepare_encode()`, `denoise_step()`, + `step_scheduler()`, or `post_decode()`, the pipeline does not implement the + required step contract. + +For implementation details and model-author guidance, see +[Diffusion Continuous Batching](../../design/feature/diffusion_continuous_batching.md). diff --git a/docs/user_guide/diffusion/request_batching.md b/docs/user_guide/diffusion/request_batching.md deleted file mode 100644 index 35bf71725..000000000 --- a/docs/user_guide/diffusion/request_batching.md +++ /dev/null @@ -1,119 +0,0 @@ -# Request-Level Batching - -Request-level batching lets diffusion serving combine multiple compatible -logical requests into one pipeline `forward()` call. Each prompt remains a -separate request with its own `request_id`, sampling parameters, seed, output, -error, and abort state. The scheduler decides which requests can run together. - -!!! warning "Prompt List Semantics" - - Diffusion request-level batching does not support a top-level packed - list-prompt request. Submit multiple prompts as independent requests and let - the scheduler batch compatible in-flight requests. Multimodal payloads stay - inside a single prompt dict, for example - `{"prompt": "...", "multi_modal_data": {"image": image}}`. - -## Enablement - -Request-level execution is the default diffusion request mode when -`step_execution=False`. `max_num_seqs` controls the maximum number of compatible -requests in one scheduler wave: - -- `max_num_seqs=1` is the serial request path. -- `max_num_seqs>1` allows fused request-level batching for pipelines that - declare request-batch support. - -Increase `max_num_seqs` above `1` to allow the request scheduler to keep more -than one compatible request active: - -```bash -vllm serve Qwen/Qwen-Image --omni \ - --port 8091 \ - --max-num-seqs 4 -``` - -For bursty online traffic, you can also set a small admission wait window. This -lets the engine wait briefly before the first `schedule()` of a new wave so -nearby compatible requests can arrive and share the same fused forward pass: - -```bash -vllm serve Qwen/Qwen-Image --omni \ - --port 8091 \ - --max-num-seqs 4 \ - --request-batch-max-wait-ms 20 -``` - -`--request-batch-max-wait-ms 0` is the default and disables admission waiting, -so there is no added wait latency. - -For deploy YAMLs, configure the diffusion stage engine args: - -```yaml -stage_args: - - stage_id: 0 - stage_type: diffusion - engine_args: - max_num_seqs: 4 - request_batch_max_wait_ms: 20 -``` - -## Compatibility - -Only pipelines that declare request-batch support can use the fused -multi-request `pipeline.forward(batch)` path. Pipelines that do not support this -contract can still use request-level execution with `max_num_seqs=1`, which is -the serial request path. Setting `max_num_seqs>1` for a pipeline without -request-batch support fails early during engine initialization. - -The scheduler batches only compatible requests. Compatibility is based on -shape-sensitive and guidance-sensitive sampling fields, including resolution, -frame count, CFG settings, output count, and LoRA identity. Requests with -different LoRA adapters or scales are kept in separate batches so the worker -activates one adapter per batch. - -Request-level batching applies when `step_execution=False`. For the separate -step-wise runtime, see [Step Execution](step_execution.md). - -## Tuning - -- `max_num_seqs` caps the number of active compatible requests in one scheduler - wave. `1` means serial request execution; values above `1` enable fused - request-level batching when the pipeline supports it. -- `request_batch_max_wait_ms` is an upper bound on extra admission wait before a - new wave starts. Keep it small for latency-sensitive serving; values such as - `10` to `50` ms are a practical starting range for bursty HTTP ingress. -- `0` disables admission waiting and preserves the lowest first-request latency. -- FIFO ordering is conservative: an incompatible request at the front of the - waiting queue can block later compatible requests from joining the current - batch. - -## Python API - -When constructing `Omni`, pass the same engine arguments: - -```python -from vllm_omni.entrypoints.omni import Omni - -omni = Omni( - model="Qwen/Qwen-Image", - max_num_seqs=4, - request_batch_max_wait_ms=20.0, -) - -outputs = omni.generate( - [ - "a cup of coffee on a table", - "a toy dinosaur on a sandy beach", - "a fox waking up in bed and yawning", - ] -) -``` - -`Omni.generate([...])` submits each list item as its own logical diffusion -request. The runtime may batch those requests internally when their sampling -parameters are compatible. - -## For Contributors - -For implementation details and model-author guidance, see -[Request-Level Batching for Diffusion](../../design/feature/diffusion_request_level_batching.md). diff --git a/docs/user_guide/diffusion/step_execution.md b/docs/user_guide/diffusion/step_execution.md deleted file mode 100644 index 4dfedfa20..000000000 --- a/docs/user_guide/diffusion/step_execution.md +++ /dev/null @@ -1,100 +0,0 @@ -# Step Execution - -Step execution is an opt-in diffusion execution mode enabled with -`step_execution=True` when constructing `Omni`. - -It is not a generic diffusion toggle for every pipeline. Only pipelines that -implement the stepwise contract support it today. - -## Quick Start - -### Python API - -```python -from vllm_omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -omni = Omni( - model="Qwen/Qwen-Image", - step_execution=True, -) - -outputs = omni.generate( - "A cat sitting on a windowsill", - OmniDiffusionSamplingParams( - num_inference_steps=50, - ), -) -``` - -### Serving - -```bash -vllm serve Qwen/Qwen-Image --omni \ - --port 8091 \ - --step-execution \ - --max-num-seqs 8 -``` - -For serving, `--step-execution` enables the step-wise runtime. Continuous -batching only becomes relevant when `--max-num-seqs > 1`. - -`max_num_seqs` controls the step scheduler capacity: - -- `max_num_seqs=1` runs one request through the step-wise state machine. -- `max_num_seqs>1` allows compatible active requests to share step waves. - -## Supported Pipelines - -| Pipeline | Example models | Step execution | -|----------|----------------|----------------| -| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | -| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No | - -!!! warning "Experimental continuous batching" - When `--step-execution` is enabled and `max_num_seqs > 1` is configured, - the step-wise path can batch - compatible requests together. This is experimental. Requests with - incompatible sampling parameters are intentionally kept in separate batches, - and `max_num_seqs=1` remains the conservative single-request step path. - -## Current Limitations - -- Continuous batching under `step_execution` is experimental and only batches - compatible requests. -- `cache_backend` is not supported together with step execution. -- Unsupported pipelines fail early during model loading. -- Request-mode extras such as KV transfer are not wired into step mode yet. -- LoRA is supported in step mode, but each batch must use a single adapter: - requests with different `lora_request` or `lora_scale` are scheduled into - separate batches. - -## When To Use It - -Use step execution only when you specifically need the pipeline to run through -its stepwise request state machine. For normal diffusion inference, leave it -disabled unless your workflow depends on this mode. - -For Qwen-Image online serving, the usual progression is: - -- start with `--step-execution --max-num-seqs 1` if you only need the step-wise path -- increase `--max-num-seqs` after that if you want the experimental compatible-request batching behavior - -If you are looking for general diffusion speedups, see -[Diffusion Features Overview](../diffusion_features.md). - -## Troubleshooting - -If model loading fails with a message mentioning `prepare_encode()`, -`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected -pipeline does not support step execution. - -## For Model Authors - -If you want to add step execution support to a new diffusion pipeline, see the -implementation guide: -[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md). - -If you also want that pipeline to participate in the experimental batched -step-wise path, see: -[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md). diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 2d9de3dba..83bc3643e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on | Method | Description | Best For | |--------|-------------|----------| -| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs | +| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs | | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss | | **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak | @@ -80,12 +80,11 @@ denoise steps. | Method | Description | Best For | |--------|-------------|----------| -| **[Request-Level Batching](diffusion/request_batching.md)** | Scheduler batches compatible independent diffusion requests into one pipeline forward pass | Bursty online serving and multi-request throughput | -| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control | +| **[Diffusion Execution Modes](diffusion/execution_modes.md)** | Configures serial requests, request batching, step execution, continuous batching, and streaming output | Matching latency, throughput, cancellation, and output-delivery requirements | -**Note:** Request-level batching is available for pipelines that declare the -request-batch forward contract. Step execution is currently supported by -QwenImagePipeline only. See [Supported Models](#supported-models) for details. +**Note:** Request-level batching and step execution are capability-based. +Consult the execution guide and the selected pipeline's documentation for +current support. ### Quantization Methods @@ -105,8 +104,8 @@ The following tables show which models support each feature: > Notes: -> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md). -> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details. +> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md). +> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details. ### ImageGen @@ -206,7 +205,7 @@ The following tables show which models support each feature: 3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible. 4. CPU Offloading (Layerwise) supports single-card for now. 5. Using FP8-Quant as an example of qunatization methods. - 6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches). + 6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches). ## Multi-Thread Weight Loading @@ -275,9 +274,9 @@ Measured on NVIDIA H800: **Memory Optimization:** -- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage +- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage - **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos -- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models +- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models **Extensions:** @@ -286,7 +285,7 @@ Measured on NVIDIA H800: **Execution Modes:** -- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support +- **[Diffusion Execution Modes](diffusion/execution_modes.md)** - Configure request batching, step execution, continuous batching, and streaming output **Startup Optimization:** diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index 7cd508d29..f1cdb9293 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -184,7 +184,7 @@ if __name__ == "__main__": !!! info For diffusion request-level batching controls such as `max_num_seqs`, see - [Request-Level Batching](../../diffusion/request_batching.md). + [Diffusion Execution Modes](../../diffusion/execution_modes.md). ### Negative Prompts diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 3a4a45573..ee36f5b60 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -5,6 +5,9 @@ path. The same configuration entrypoint is used across diffusion-only models, multi-stage omni/TTS models, and multi-stage diffusion models, but each model type has a different quantization scope. +For the internal architecture and backend extension points, see the +[quantization design document](../../design/feature/quantization.md). + ## Quantization Modes | Mode | Guide | Description | Methods | diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 985c4ff7d..6ebc7ae63 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -260,7 +260,7 @@ if __name__ == "__main__": For diffusion pipelines, the input list is sliced into single-item requests before feeding into the diffusion pipeline. For request-level batching controls such as `max_num_seqs`, see - [Request-Level Batching](../../../docs/user_guide/diffusion/request_batching.md). + [Diffusion Execution Modes](../../../docs/user_guide/diffusion/execution_modes.md). ### Negative Prompts diff --git a/mkdocs.yml b/mkdocs.yml index 8d5ffaa1b..fd97b9b37 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -71,6 +71,14 @@ plugins: - search - autorefs - awesome-nav + - redirects: + redirect_maps: + user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md + user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md + user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md + user_guide/quantization/index.md: user_guide/quantization/overview.md + design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md + design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md - glightbox - git-revision-date-localized: # exclude files diff --git a/pyproject.toml b/pyproject.toml index 4af374589..f8328e656 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,6 +108,7 @@ docs = [ "mkdocstrings-python", "mkdocs-gen-files", "mkdocs-awesome-nav", + "mkdocs-redirects", "mkdocs-glightbox", "mkdocs-git-revision-date-localized-plugin", "mkdocs-minify-plugin", @@ -249,13 +250,16 @@ markers = [ "npu: Tests that run on NPU/Ascend (auto-added)", "musa: Tests that run on MUSA/Moore Threads (auto-added)", # specified computation resources marks (auto-added) - "H100: Tests that require H100 GPU", - "L4: Tests that require L4 GPU", - "MI325: Tests that require MI325 GPU (AMD/ROCm)", - "B60: Tests that require Intel Arc Pro B60 XPU", - "S5000: Tests that require S5000 GPU (Moore Threads/MUSA)", - "A2: Tests that require A2 NPU", - "A3: Tests that require A3 NPU", + # Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list(). + "H100: [hardware-resource] Tests that require H100 GPU", + "H200: [hardware-resource] Tests that require H200 GPU", + "L4: [hardware-resource] Tests that require L4 GPU", + "B200: [hardware-resource] Tests that require B200 GPU", + "MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)", + "B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU", + "S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)", + "A2: [hardware-resource] Tests that require A2 NPU", + "A3: [hardware-resource] Tests that require A3 NPU", "distributed_cuda: Tests that require multi cards on CUDA platform", "distributed_rocm: Tests that require multi cards on ROCm platform", "distributed_xpu: Tests that require multi cards on XPU platform", diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index fa9ec6755..3142f9fd1 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -332,6 +332,6 @@ bundled offline CLI do not currently expose `sigmas`. - - - [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md) -- [Request batching](../../docs/user_guide/diffusion/request_batching.md) +- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md) - [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md) - [I2V offline example](../../examples/offline_inference/image_to_video/README.md) diff --git a/recipes/MiniMaxAI/MiniMax-H3-5090.md b/recipes/MiniMaxAI/MiniMax-H3-5090.md new file mode 100644 index 000000000..fb8b0f118 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-5090.md @@ -0,0 +1,61 @@ +# MiniMax-H3 on RTX 5090 + +This recipe uses BF16 weights, tiled VAE decode, tensor parallelism where a +second GPU is available, and distributed layerwise offload (DLO). It is a +memory-first serving configuration; lower resident counts reduce HBM use and +increase CPU-to-GPU transfer time. + +## Capacity requirements + +| Resource | One RTX 5090 | Two RTX 5090s | +| --- | ---: | ---: | +| GPU HBM | 32 GiB | 32 GiB per GPU | +| Checkpoint storage | 135 GiB per partition | 135 GiB per partition | +| Available system RAM | 200 GiB minimum | 200 GiB minimum | +| Recommended system RAM | 384 GiB | 384 GiB | + +`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one +server at a time. DLO keeps rank-local weights in pinned host memory; increasing +`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the +current implementation because resident layers retain pinned CPU master copies. + +> **Modular H3:** after #5720 lands, preserve this recipe's one-partition +> behavior with `--task-type fl2va` or `--task-type ref2va`. + +## One RTX 5090: 1344x768, 5 seconds + +Use 12 resident DiT layers. A 50-step B300 allocation test with this exact +single-rank topology peaked at 26.50 GiB; re-measure peak HBM on the target +card before increasing the resident count. + +```bash +CUDA_VISIBLE_DEVICES=0 vllm serve /path/to/MiniMax-H3/FL2VA \ + --omni --trust-remote-code --host 0.0.0.0 --port 8000 \ + --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \ + --usp 1 --ring 1 --vae-patch-parallel-size 1 \ + --vae-parallel-mode tile --vae-use-tiling \ + --enable-distributed-layerwise-offload --dlo-no-use-allgather \ + --dlo-resident-layers 12 --enforce-eager \ + --diffusion-attention-backend CUDNN_ATTN +``` + +## Two RTX 5090s: 1344x768, 5 seconds + +Use TP2 and 20 resident DiT layers. The two-rank B300 capacity run peaked at +27,726 MiB per rank for this shape and 50 steps. This is a memory/correctness +proxy, not a consumer-GPU latency claim. + +```bash +CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \ + --omni --trust-remote-code --host 0.0.0.0 --port 8000 \ + --num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \ + --usp 1 --ring 1 --vae-patch-parallel-size 2 \ + --vae-parallel-mode tile --vae-use-tiling \ + --enable-distributed-layerwise-offload --dlo-no-use-allgather \ + --dlo-resident-layers 20 --enforce-eager \ + --diffusion-attention-backend CUDNN_ATTN +``` + +For Ref2VA, stop the FL2VA server and restart the same command with +`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video count and prompt length +can increase activation memory; begin with one request at a time. diff --git a/recipes/MiniMaxAI/MiniMax-H3-MUSA.md b/recipes/MiniMaxAI/MiniMax-H3-MUSA.md new file mode 100644 index 000000000..8d5053759 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-MUSA.md @@ -0,0 +1,146 @@ +# MiniMax H3 on Moore Threads MUSA + +> Joint video and audio generation on Moore Threads GPUs + +## Summary + +- Vendor: MiniMaxAI +- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) +- Tasks: T2VA, FL2VA, and Ref2VA +- Mode: OpenAI-compatible `/v1/videos` HTTP serving +- Hardware: 4x Moore Threads MTT S5000 for the validated Ref2VA profile +- Maintainer: Community + +This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for MUSA environments. +The modular H3 pipeline can load both task-specific DiTs, but this MUSA recipe +uses a **single task partition**: `FL2VA` for `t2va`/`fl2va`, or `Ref2VA` for +`ref2va`. This avoids loading both DiTs. + +## Prerequisites + +### Checkpoint + +Download the partition required by the selected startup task. The commands below +run a single-task server; do not use the modular default (combined) mode on MUSA +until its two-DiT memory and performance profile is validated. + +For T2VA and FL2VA: + +```bash +python -m pip install modelscope +export MODEL_ROOT=/path/to/MiniMax-H3 +modelscope download MiniMax/MiniMax-H3 \ + --local_dir "${MODEL_ROOT}" \ + --max-workers 16 \ + --include 'FL2VA/**' +``` + +For Ref2VA, download `Ref2VA/**` instead. + +### Environment + +Install a compatible PyTorch/torch-musa, torchada, MATE, and vLLM-MUSA +stack before installing vLLM-Omni. Importing vLLM-Omni should report that the +`musa` vLLM platform plugin is active. + +Install vLLM-Omni from a checkout containing MiniMax H3 support: + +```bash +python -m pip install -e . +``` + +MiniMax H3 reference inputs and MP4 output require `soundfile`, `ffmpeg`, and +`ffprobe`. The Python dependency is installed by vLLM-Omni; install the two +executables with the operating system package manager and verify them: + +```bash +python -c 'import soundfile; print(soundfile.__version__)' +ffmpeg -version +ffprobe -version +``` + +When torchaudio cannot load TorchCodec, vLLM-Omni automatically falls back to +soundfile. Formats that libsndfile cannot read are demuxed through ffmpeg. + +## Start a server + +Use the partition path and set `--task-type` explicitly: `fl2va` for T2VA/FL2VA +or `ref2va` for Ref2VA. + +The validated Ref2VA configuration uses tensor parallelism across four MTT S5000 GPUs and offloads inactive model components to CPU: + +```bash +export MODEL="${MODEL_ROOT}/Ref2VA" +export PORT=8091 + +MUSA_VISIBLE_DEVICES=0,1,2,3 \ +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --num-gpus 4 \ + --task-type ref2va \ + --tensor-parallel-size 4 \ + --vae-patch-parallel-size 4 \ + --enable-cpu-offload \ + --vae-use-tiling \ + --diffusion-attention-backend FLASH_ATTN +``` + +On MUSA, `FLASH_ATTN` selects the MATE/FlashAttention-3 implementation. Do +not install the CUDA-only `fa4` optional dependency. + +For request examples and parameter definitions, see +[MiniMax-H3.md](MiniMax-H3.md#http-api-examples). + +## MUSA validation + +The following weight-independent gates have been validated on one MTT S5000: + +- packed variable-length attention selects MATE FlashAttention-3 and returns + finite BF16 output; +- MiniMax H3 RoPE casts CPU FP64 position metadata to FP32 before device + arithmetic and returns finite output; +- the video-VAE seeded RNG context is deterministic on MUSA and restores both + CPU and MUSA RNG states; +- the soundfile fallback preserves stereo channel layout, FP32 samples, and + the native sample rate when torchaudio/TorchCodec is unavailable. + +Ref2VA image+audio generation was also validated end to end on one MTT S5000 +with CPU offload and VAE tiling. The smoke used a 448x256 reference image, a +four-second 32 kHz stereo audio reference, two denoising steps, and seed 42. +The synchronous HTTP request completed in 50.6 seconds and returned an MP4 +with the following decoded properties: + +- 107 video frames at 448x256 and 24 FPS; +- non-static video (`pixel_std=13.07`, `temporal_delta=0.49`); +- finite stereo audio at 32 kHz (`audio_rms=0.0458`). + +Official Ref2VA video-reference serving was validated with the expanded +MiniMax prompt at 1344x768, 5 seconds, seed 0, and 50 steps. A true TP4 run on +four MTT S5000 GPUs returned HTTP 200 in 638.6 seconds and produced 124 H.264 +frames at 24 FPS with AAC stereo audio at 32 kHz. + +Keep `--vae-use-tiling` enabled for this serving profile. + +## Known limitations + +- Keep Ring Attention at degree 1. The current Ring path does not preserve + MiniMax H3 packed padding boundaries. +- H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. +- This MUSA recipe uses an explicitly selected single task partition; combined + serving loads both DiTs. +- H3 currently executes one generation request per diffusion batch. +- FP8 quantization has not been enabled for MiniMax H3. +- MP3, M4A, MP4, and reference-video audio fallback requires `ffmpeg` on + `PATH`; WAV inputs can be read directly through soundfile. + +## Additional resources + +- [MiniMax-H3.md](MiniMax-H3.md) — full model and API guide +- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — Ascend NPU deployment guide +- [Supported models](../../docs/models/supported_models.md) +- [Video API](../../docs/serving/videos_api.md) diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md index e6120e916..930411f1a 100755 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md @@ -26,12 +26,12 @@ environments. Differences from the GPU path: ### Checkpoint -Same as the GPU recipe — Hugging Face access approval is required: +Same as the GPU recipe — Hugging Face access approval is required. Authenticate +once; `vllm serve` downloads the required nested components automatically: ```bash hf auth login -export MODEL_ROOT=/path/to/MiniMax-H3 -hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}" +export MODEL=MiniMaxAI/MiniMax-H3 ``` ### Environment @@ -50,7 +50,8 @@ uv pip install -e . ``` Install the **mindie-sd** operator library to enable Ascend-optimized fused -operators (`adalayernorm`, etc.): +operators (`adalayernorm`, etc.) and the RainFusion `rf_v2` block-sparse +attention kernel: ```bash git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD @@ -72,24 +73,21 @@ pip install mindiesd-*.whl ## Start a server -One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA -and FL2VA requests, or to `Ref2VA` for Ref2VA requests. +Pass the repository ID directly. The pipeline loads the two nested DiTs while +sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`. -### Multi-NPU: 768P validated configuration +### Multi-NPU: 768P combined-service configuration -Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence -parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled -VAE patch parallelism degree 8, and layerwise offload: +Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism +degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload: ```bash export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 export PORT=9098 -export MODEL="${MODEL_ROOT}/FL2VA" +export MODEL=MiniMaxAI/MiniMax-H3 export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 export PYTHONDONTWRITEBYTECODE=1 -export HF_HUB_OFFLINE=1 -export TRANSFORMERS_OFFLINE=1 vllm serve "${MODEL}" \ --omni \ @@ -111,12 +109,31 @@ Do not add `--enforce-eager`. The first request includes regional compilation; warm the server once before measuring steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. -To serve Ref2VA, stop the FL2VA server and restart with: +The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no +partition restart is required. Layerwise offload applies to both DiTs. + +### Optional optimizations + +Two independent optimizations may be enabled on top of the configuration +above. Both are validated for T2VA only. + +**RainFusion block-sparse attention** — switch the attention backend, keeping +every other flag unchanged: ```bash -export MODEL="${MODEL_ROOT}/Ref2VA" + --diffusion-attention-backend RAINFUSION_ATTN ``` +**INT8 online quantization** — add one flag to the server command above: + +```bash + --quantization int8 +``` + +Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks +over the whole sequence, so ring parallelism would split away the keys it +needs. Scale with `--usp` instead. + ## HTTP API examples The request format is identical to the GPU recipe; see @@ -145,7 +162,8 @@ throughput guarantee. ## Known limitations -- Each server process loads only one checkpoint partition. +- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads + both task-specific DiTs, and loads one copy of every shared component. - H3 currently executes one generation request per diffusion batch. - The first regional-compile request is a warmup and should not be included in steady-state performance measurements. @@ -154,9 +172,16 @@ throughput guarantee. the same HTTP request contract. - VAE patch parallelism requires size 1 or the full DiT group size and supports the H3 native `tile` mode only. +- RainFusion block-sparse attention and INT8 quantization are validated for + T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA. +- Online quantization cannot be combined with distributed layerwise offload + while AllGather is enabled; pass `--dlo-no-use-allgather` in that case. ## Additional resources - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide +- [Attention backends § RAINFUSION_ATTN](../../docs/user_guide/diffusion/attention_backends.md#rainfusion_attn-backend-and-block-sparse-video-attention) + — RainFusion knobs and tuning +- [Int8 quantization](../../docs/user_guide/quantization/int8.md) - [Supported models](../../docs/models/supported_models.md) - [Video API](../../docs/serving/videos_api.md) diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index dbdc34ad5..cfe9b5796 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -12,33 +12,49 @@ - Maintainer: Community MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its -checkpoint has two independently served partitions: +checkpoint has two task-specific DiT partitions: - `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio (`fl2va`) - `Ref2VA`: up to 9 images, 3 videos, and 3 audio references in supported image/video/audio combinations (`ref2va`; audio-only is rejected) +One vLLM-Omni diffusion stage can load both DiTs while instantiating the +tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only +once. Requests select the DiT with `extra_params.task`. + The generated MP4 contains H.264 video and synchronized stereo audio. ## Prerequisites -The checkpoint requires Hugging Face access approval. Download it and point -`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA` -subdirectories: +The checkpoint requires Hugging Face access approval. Authenticate once; +`vllm serve` downloads the required components automatically: ```bash hf auth login -export MODEL_ROOT=/path/to/MiniMax-H3 -hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}" +export MODEL=MiniMaxAI/MiniMax-H3 ``` -Install vLLM-Omni from the checkout containing MiniMax H3 support. On -Blackwell, install the optional FlashAttention-4 dependency: +The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and +`Ref2VA/transformer/**`. It does not download or load the diffusers-format +`transformer`, `transformer_ref`, or `vae` weights at the repository root, nor +duplicate Ref2VA copies of shared components. + +Install vLLM-Omni from the checkout containing MiniMax H3 support. The +two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need +FlashAttention-4. Install the optional dependency only for the four-GPU +B300/GB200 `FLASH_ATTN` profile: ```bash uv venv source .venv/bin/activate +uv pip install -e . +``` + +To keep FA4 available as an explicit option on Blackwell, install the +FlashAttention-4 extra: + +```bash uv pip install -e '.[fa4]' ``` @@ -47,8 +63,29 @@ reference-video preparation and MP4 output. ## Start a server -One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA -and FL2VA requests, or to `Ref2VA` for either Ref2VA request. +Pass the repository ID directly. The pipeline uses `FL2VA` for model discovery +and shared components, and loads the second DiT from `Ref2VA/transformer`. + +### Memory and storage requirements + +Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each +H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of +BF16 safetensors (about **135 GiB** on disk). Keeping both partitions +locally therefore needs roughly **270 GiB** of model storage. A combined +service downloads both; `--task-type fl2va` or `--task-type ref2va` downloads +only the selected partition. + +CPU offload and distributed layerwise offload reduce GPU residency; they do +not make the model weights disappear. With `--dlo-no-use-allgather`, each +worker retains its standard-loader rank-local weights in host memory, including +pinned CPU buffers used for H2D streaming. Use at least **200 GiB available +system RAM** before starting the two-GPU recipe; a **384 GiB host is +recommended** to leave room for the OS, CUDA/PyTorch allocations, request +inputs, and filesystem cache. Do not run the FL2VA and Ref2VA servers at the +same time on a host sized for this minimum. + +The consumer-GPU profiles below are HBM budgets only. They still require the +host-RAM budget above. ### Single GPU: accuracy and memory first @@ -57,7 +94,7 @@ This matches the accuracy-qualified reference path and prevents the Qwen3-VL encoder and DiT from being resident on the GPU at the same time. ```bash -export MODEL="${MODEL_ROOT}/FL2VA" +export MODEL=MiniMaxAI/MiniMax-H3 export PORT=8091 CUDA_VISIBLE_DEVICES=0 \ @@ -74,26 +111,120 @@ vllm serve "${MODEL}" \ ``` Use a GPU with enough memory for the active H3 component and enough system RAM -for the offloaded components. CPU offload reduces GPU memory pressure but adds -PCIe/NVLink transfer latency. +for both offloaded DiTs plus the shared components. Model-level offload keeps +the two DiTs mutually exclusive on GPU, but adds PCIe/NVLink transfer latency. -### Four GPUs: measured best-practice throughput +### Two 24/32 GB GPUs: TP2 distributed layerwise offload -The best validated four-GPU configuration on four NVIDIA B300 GPUs is: +For two PCIe consumer GPUs, combine TP2 with distributed layerwise offload +(DLO). The standard loader first creates the rank-local TP shard. DLO keeps +that shard in pinned host memory and streams the 30 tail DiT blocks through a +shared two-buffer window without a DP AllGather. The first 20 DiT blocks are +copied to the GPUs once per denoise stage, reused by every sampling step, and +released before VAE decode so the decoder can reuse their HBM. + +```bash +export MODEL=MiniMaxAI/MiniMax-H3 +export PORT=8091 + +CUDA_VISIBLE_DEVICES=0,1 \ +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --task-type fl2va \ + --num-gpus 2 \ + --tensor-parallel-size 2 \ + --usp 1 \ + --ring 1 \ + --text-encoder-tp-size 2 \ + --vae-patch-parallel-size 2 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --dlo-resident-layers 20 \ + --enforce-eager \ + --diffusion-attention-backend CUDNN_ATTN +``` + +Use the profile that matches the per-GPU memory capacity: + +| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status | +|---|---:|---:|---:|---|---|---| +| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated | +| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point | + +This topology uses all available parallel capacity: TP2 shards both the DiT +and text encoder, `--dlo-no-use-allgather` streams each rank's local TP shard +without reconstructing full blocks, and VAE patch parallelism splits tiled +decode across both GPUs. cuDNN attention is selected explicitly for the RTX +consumer path; the server stays eager to avoid an unqualified compile path. + +The resident count changes placement and transfer frequency only; it does not +quantize or change the BF16/FP32 denoise math. Re-measure peak memory before +increasing it on a different request shape. + +### RTX 5090 target-hardware validation + +At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on +2 x RTX 5090 without OOM: + +| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation | +|---:|---:|---:|---:|---| +| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed | + +This is a single end-to-end validation run, not a warmed multi-run latency +benchmark. The sampled `nvidia-smi` peak is also not a CUDA allocator +high-water mark. The environment used vLLM 0.26.0, vLLM-Omni +`0.26.1.dev14+gae6577ea`, and PyTorch 2.11.0+cu130. The +[run record](https://github.com/lishunyang12/vllm-omni-rankings/blob/dcd06d7e83cb069842535918c0169ee9f3f29ba0/scripts/%E5%BE%AE%E4%BF%A1%E5%9B%BE%E7%89%87_20260805000034_86_237.png) +captures the environment, output contract, elapsed time, and sampled peak. + +Before the target run, both profiles were exercised on two B300 ranks as an +allocation and correctness proxy. At 1344x768, 124 frames, and 50 steps, the +20-layer profile peaked at 27,726 MiB per rank. At 1024x576, the 12-layer +profile peaked at 18,888 MiB per rank in a 5-step capacity run. The resident +and fully streamed placements produced identical decoded video-frame and audio +hashes for the same shape, step count, prompt, and seed. The B300 result does +not establish RTX 4090 PCIe latency; treat the 4090 profile as a conservative +starting point until it is measured on that GPU. + +To run T2VA, FL2VA, image+audio Ref2VA, and two-video Ref2VA in order, validate +every MP4's H.264/AAC streams, and retain live server and GPU-memory logs: + +```bash +RUN_ROOT=/path/to/run-root \ +MODEL_ROOT=/path/to/MiniMax-H3 \ +GPU_IDS=0,1 \ +PROFILE=rtx5090 \ +bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh +``` + +The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for +`PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default. + +### Four GPUs: throughput-oriented combined service + +For a combined service on four high-memory GPUs, use: - no CPU or layerwise offload; - Ulysses sequence parallelism degree 4; - native tiled VAE patch parallelism degree 4; - regional `torch.compile` for the repeated DiT blocks; -- CuTe FlashAttention-4 through the `FLASH_ATTN` backend, with Ring and TP - left at 1. +- dense BF16 `TRTLLM_ATTN`, with Ring and TP left at 1. + +Both DiTs remain resident in this no-offload configuration. If they do not fit, +use model-level CPU offload. ```bash -export MODEL="${MODEL_ROOT}/FL2VA" +export MODEL=MiniMaxAI/MiniMax-H3 export PORT=8091 CUDA_VISIBLE_DEVICES=0,1,2,3 \ -FLASHINFER_DISABLE_VERSION_CHECK=1 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ vllm serve "${MODEL}" \ @@ -106,19 +237,76 @@ vllm serve "${MODEL}" \ --ring 1 \ --vae-patch-parallel-size 4 \ --vae-parallel-mode tile \ - --vae-use-tiling \ - --diffusion-attention-backend FLASH_ATTN + --vae-use-tiling ``` -On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains -`Using CuTe FlashAttention-4 on Blackwell` before recording measurements. - Do not add `--enforce-eager` to this performance configuration. The first request includes regional compilation; warm the server once before measuring steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. The H3 VAE supports its native `tile` mode, not `spatial_shard_height` or `spatial_shard_width`. +### Attention Backends + +On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16 +`TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly, +use: + +```bash +--diffusion-attention-backend TRTLLM_ATTN +``` + +Stable measurements with the four-GPU profile above put dense `TRTLLM_ATTN` +and FA4 within 2% of each other. `TRTLLM_ATTN` remains the datacenter Blackwell +default and enables the optional optimizations below. Confirm the server log +contains `Defaulting to diffusion attention backend TRTLLM_ATTN` before +recording measurements when using the default selection. + +FA4 remains available by explicitly selecting the `FLASH_ATTN` backend: + +```bash +--diffusion-attention-backend FLASH_ATTN +``` + +On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains +`Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements. + +`TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main +DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse +Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses +`fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V +to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short +attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled +for it. The example enables the calibration-free Skip-Softmax path with +`threshold=0.05`, after the normalized timestep reaches `0.97`: + +```bash +--diffusion-attention-config '{ + "default": { + "backend": "TRTLLM_ATTN", + "quant": { + "dtype_qk": "fp8_e4m3", + "q_block_size": 1, + "k_block_size": 16 + }, + "skip_softmax": { + "threshold": 0.05, + "disabled_until_timestep": 0.97 + } + }, + "per_role": { + "minimax_h3.token_refiner": { + "backend": "TRTLLM_ATTN" + } + } +}' +``` + +For configuration details, see +[TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax) +and +[TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization). + ### Text encoder tensor parallelism The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by @@ -129,11 +317,10 @@ vLLM-style tensor-parallel layers and runs with distributed collectives over its own encoder process group): ```bash -export MODEL="${MODEL_ROOT}/FL2VA" +export MODEL=MiniMaxAI/MiniMax-H3 export PORT=8091 CUDA_VISIBLE_DEVICES=0,1,2,3 \ -FLASHINFER_DISABLE_VERSION_CHECK=1 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ vllm serve "${MODEL}" \ @@ -147,8 +334,7 @@ vllm serve "${MODEL}" \ --text-encoder-tp-size 4 \ --vae-patch-parallel-size 4 \ --vae-parallel-mode tile \ - --vae-use-tiling \ - --diffusion-attention-backend FLASH_ATTN + --vae-use-tiling ``` `N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so @@ -161,12 +347,9 @@ the reference path within bf16 rounding: every encoder rank all-reduces the row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is replicated on every rank. -To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU -command with: - -```bash -export MODEL="${MODEL_ROOT}/Ref2VA" -``` +No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while +`task=ref2va` routes to +`Ref2VA/transformer`. T2VA uses the FL2VA DiT. ### Online FP8 quantization @@ -222,7 +405,7 @@ export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync" ### 1. T2VA: text to video and audio -Run this request against the `FL2VA` partition: +Run this request against the combined service: ```bash curl -sS -X POST "${API_URL}" \ @@ -240,7 +423,7 @@ curl -sS -X POST "${API_URL}" \ ### 2. FL2VA: first frame to video and audio -Run this request against the `FL2VA` partition. When width and height are +Run this request against the combined service. When width and height are omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short edge. @@ -279,10 +462,11 @@ curl -sS -X POST "${API_URL}" \ ### 3. Ref2VA: image-only, image/audio, or mixed references -Run these requests against the `Ref2VA` partition. Image-only Ref2VA omits -`audio_reference`; adding one or more audio references is optional. The typed -fields accept one object or an ordered JSON list. `audio_reference` accepts an -HTTP(S) URL or a `data:` URL. +Run these requests against the combined service or a Ref2VA-only service. +Image-only Ref2VA omits `audio_reference`; adding one or more audio references +is optional. The typed fields accept one object or an ordered JSON list. +`audio_reference` accepts an HTTP(S) URL or a `data:` URL. In one terminal, +expose the local reference assets to the serving host: ```bash python -m http.server 8092 \ @@ -332,7 +516,7 @@ shorter, the reference soundtrack is truncated to the generated clip. ### 4. Ref2VA: video, separate audio, and mixed references -Run this request against the `Ref2VA` partition. Repeat the +Run this request against the combined service. Repeat the `input_references` multipart field once per source video. H3 consumes the videos in form order and preserves their original soundtracks during conditioning. @@ -394,7 +578,7 @@ than one is requested. | Parameter | Recommended value | Notes | |-----------|-------------------|-------| -| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition | +| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT | | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count | | `fps` | `24` | H3 output FPS is fixed | | `num_inference_steps` | `50` | Matches the reference accuracy workloads | @@ -407,6 +591,10 @@ than one is requested. | `start_time_seconds` | `0` | Reference-video segment start; use a list in `extra_params` for multiple videos | | `width`, `height` | Multiples of 32 | Output aspect ratio must be between 1:4 and 4:1 | +## ComfyUI Frontend + +Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 service. The ComfyUI frontend can run in a separate environment or machine. Refer to [vLLM-Omni ComfyUI Integration](../../docs/features/comfyui.md) for details. + ## Validated four-GPU evidence The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one @@ -432,7 +620,8 @@ reduction. ## Known limitations -- Each server process loads only one checkpoint partition. +- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads + both task-specific DiTs, and loads shared components once from `FL2VA`. - H3 currently executes one generation request per diffusion batch. - The first regional-compile request is a warmup and should not be included in steady-state performance measurements. diff --git a/recipes/README.md b/recipes/README.md index 507290fb5..a6015d9b0 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -45,7 +45,8 @@ recipes/ | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) | | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | -| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, first/last-keyframe FL2VA, and full mixed-reference Ref2VA serving | 1x GPU with CPU offload / 4x B300 | +| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 | +| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) | | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU | | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index f169e6a6b..a9a63d5ce 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -164,10 +164,11 @@ curl -s http://localhost:8091/v1/chat/completions \ PR [#2495](https://github.com/vllm-project/vllm-omni/pull/2495) adds performance CI configs for the same DiT-only settings. The CI step is currently opt-in (gated by `RUN_HUNYUAN_IMAGE3_PERF=1`) with `soft_fail` -enabled, intended for initial data collection. Performance assertions are -skipped (`skip-performance-assertion: true`); the baseline values in the -JSON configs are reference-only and will be promoted to regression gates -once enough nightly data has been collected. +enabled, intended for initial data collection. Perf runners no longer +compare metrics against JSON baselines at assert time; the `baseline` +blocks in the configs are reference-only and are copied into result JSON +for reporting. Regression gates can be reintroduced once enough nightly +data has been collected. The user-facing equivalent is to launch one of the CLI commands above and generate 1024x1024 images with 50 denoising steps. diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md index 27d1c77bd..8eeea5cc2 100644 --- a/recipes/cosmos3/Cosmos3-DistOffload.md +++ b/recipes/cosmos3/Cosmos3-DistOffload.md @@ -56,10 +56,14 @@ vllm serve /path/to/Cosmos3-Super \ --data-parallel-size 4 ``` -### Without AllGather (full weights per rank, no sharding) +### Without DLO AllGather (rank-local weights, no DLO sharding) -Use `--dlo-no-use-allgather` when you want each rank to load full weights independently — -no AllGather synchronization, no concurrent request requirement, but N× host memory. +Use `--dlo-no-use-allgather` when you want each rank to stream the standard +loader's rank-local weights independently — no DLO AllGather synchronization, +no concurrent request requirement, but no DLO weight sharding and higher host +memory use. In pure DP, this is equivalent to independent full-weight +layerwise-offload replicas; with TP or HSDP, each rank retains the standard +loader's local tensor layout. ```bash vllm serve /path/to/Cosmos3-Nano \ @@ -154,9 +158,9 @@ wait | Flag | Description | Default | |------|-------------|---------| | `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` | -| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` | +| `--data-parallel-size N` | Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled | `1` | | `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` | -| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` | +| `--dlo-no-use-allgather` | Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather) | `false` | | `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` | ## Mode comparison @@ -166,7 +170,7 @@ wait | DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences | | DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory | | DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) | -| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU | +| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Standard-loader rank-local weights | 2 blocks | 0.4× HSDP | No DLO AllGather overhead, high CPU | | HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM | ## Memory expectations @@ -185,7 +189,11 @@ 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 quantization (FP8) is incompatible with mmap loading — falls back to regular - `load_weights()` automatically -- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension) -- HSDP + AllGather is rejected (would double-shard weights) +- 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 +- HSDP + AllGather is rejected (would double-shard weights). HSDP with + `--dlo-no-use-allgather` is accepted at configuration level, but is not yet a + broadly validated production combination diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index c3282ba51..b27d1b473 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -423,6 +423,156 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_ once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`; unknown keys for the model are dropped. +## ROCm + +### 1x MI350X (Online serving) + +#### Environment + +- OS: Ubuntu 22.04+ +- Python: 3.12+ +- Driver / runtime: ROCm, HIP 7.2, **gfx942 or gfx950** (validated on gfx950 / MI350X, + which exposes 252 GiB of HBM to PyTorch) +- PyTorch: 2.11.0 (ROCm build) +- vLLM version: match the repository requirements from your current checkout +- vLLM-Omni version or commit: use the commit you are deploying from + +No ROCm-specific patch or flag is needed: Cosmos3 runs on the stock code path. + +#### Command + +The NVIDIA command works unchanged. Guardrails need the gated +`nvidia/Cosmos-1.0-Guardrail` repo, so the quick path is to disable them (you are +then responsible for license compliance): + +```bash +vllm serve nvidia/Cosmos3-Nano \ + --omni \ + --no-guardrails \ + --host 0.0.0.0 --port 8000 \ + --init-timeout 1800 +``` + +For 720p video, add `--vae-use-tiling`: it cuts peak VRAM by ~68% for ~13% +latency, which is the difference between needing a 120 GiB card and a 40 GiB one. +`--quantization fp8` and `--enable-layerwise-offload` are also supported; see the +Notes for the measured cost of each. For extra GPUs use `--ulysses-degree N` or +`--tensor-parallel-size N` (not validated on ROCm yet). + +#### Verification + +Use the same `curl` requests as the CUDA section above. The ROCm images do not +ship `ffprobe`, so the CUDA section's `ffprobe` line does not work here. `imageio` +and `pyav` are already installed and report the same thing: + +```bash +python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print(im.size, im.mode)" +python -c "import imageio.v3 as iio; print(iio.improps('cosmos3_t2v.mp4', plugin='pyav').shape)" +``` + +Expect `(1024, 1024) RGB` and `(189, 720, 1280, 3)` for the requests above. +Re-running the same request with the same seed reproduces a byte-identical mp4, +so an md5 comparison across two runs also works as a smoke check. + +#### Notes + +- **Warm up once per output shape before timing anything.** The first generation + of a shape pays a one-time cost that is cached on disk and reused by later + processes in the same container: a ~250 s aiter attention-kernel JIT build plus + a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images + does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is + **540 s cold vs 161 s warm, with byte-identical output**. +- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps + **~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is + the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the + denoise loop. The optional flags act on different terms of the memory bill and + are therefore complementary rather than redundant: + + | Flag | Latency | Peak reserved | Peak allocated | Acts on | + |---|---|---|---|---| + | *(none)* | 161 s | 120 GiB | 95 GiB | — | + | `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations | + | `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights | + | `--quantization fp8` | 150 s (−6.6%) | 107 GiB (−11%) | 82 GiB (−14%) | weights (online, no calibration) | + + Reserved is the caching-allocator high water mark and runs ~2–25 GiB above + allocated; allocated is the figure to compare against another platform. +- **`--vae-use-tiling` is by far the biggest lever for video.** Of the 95 GiB + allocated by default only ~35 GiB is weights and fixed overhead; the remaining + **60 GiB is activations from decoding all 189 frames in one piece**, which is + why fp8 and offload cannot touch it and why tiling, which leaves under 5 GiB of + them, wins so much. `--vae-use-slicing` is a no-op for a single video request: + peak memory is bit-identical to the default. +- **`--enable-layerwise-offload` is close to free on this part:** per-block DiT + compute is long enough to hide the weight transfer, so on a card this size it + buys headroom for concurrency or higher resolution rather than the + speed-for-memory trade it is on small cards. +- **Attention backend:** the numbers above are aiter FlashAttention, not SDPA. + ROCm picks aiter when it is installed and the encoded device capability + satisfies `90 < major * 10 + minor < 100` (gfx942 and gfx950); otherwise it + falls back to `TORCH_SDPA` and says so only at debug level, which is easy to + miss. If you are reporting numbers, assert the choice instead of trusting the + default: `RocmOmniPlatform.get_diffusion_attn_backend_cls(None, 128)` returns + the resolved path. +- **Determinism is per-configuration:** same seed and same flags reproduce a + byte-identical mp4 across processes, but different configurations do not match + each other — including offload, which nominally only relocates weights, so do + not use it as a bit-exact regression baseline. Quality was not evaluated for + any configuration. + +### 1x MI350X (Offline generation) + +#### Environment + +Same as the online serving section above. + +#### Command + +The standard task examples run unchanged: + +```bash +# Text-to-image +python examples/offline_inference/text_to_image/text_to_image.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A photorealistic red sports car at golden hour, cinematic lighting." \ + --negative-prompt "blurry, distorted, low quality" \ + --height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \ + --extra-body '{"flow_shift": 3.0, "guardrails": false}' \ + --output cosmos3_t2i.png + +# Text-to-video +python examples/offline_inference/text_to_video/text_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A robot arm is cleaning a plate in the kitchen." \ + --negative-prompt "blurry, distorted, low quality" \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 --flow-shift 10.0 \ + --extra-body '{"max_sequence_length": 4096, "guardrails": false, + "use_resolution_template": false, "use_duration_template": false}' \ + --output cosmos3_t2v.mp4 +``` + +Both scripts print `Total generation time` and `Worker peak GPU memory +(reserved)`, which is where the latency and reserved figures in the Notes above +come from. Add `--vae-use-tiling`, `--quantization fp8` or +`--enable-layerwise-offload` to reproduce the other rows. The allocated column +needs `VLLM_LOGGING_LEVEL=DEBUG`, which enables a per-request line reporting +reserved, allocated and pool overhead together. + +#### Verification + +Same as the online serving section above. + +#### Notes + +- Run each command twice and read the second run; the first pays the one-time + compilation described above. If you run in a container, keep the container + alive between runs (`docker run --rm` throws the JIT cache away and you pay + the ~250 s kernel build every time). +- Latency and memory figures are identical to the online serving section — they + were taken from this path, because a single process has no client/server + timing ambiguity. + ## NPU ### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 6a1b263d2..c0a3f20e8 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -39,7 +39,7 @@ from vllm_omni.config.stage_config import ( load_deploy_config, merge_pipeline_deploy, ) -from vllm_omni.engine.stage_init_utils import build_engine_args_dict +from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -124,10 +124,16 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False) assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto") assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto") - assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization", 0.90) - assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs", 128) + assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization") + assert omni_stage.cache_config.enable_prefix_caching == engine_args.get("enable_prefix_caching") + expected_disable_hybrid = engine_args.get("disable_hybrid_kv_cache_manager") + if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_GENERATION: + expected_disable_hybrid = True if expected_disable_hybrid is None else expected_disable_hybrid + assert omni_stage.cache_config.disable_hybrid_kv_cache_manager == expected_disable_hybrid + assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs") assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens") - assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling", True) + assert omni_stage.scheduler_config.enable_chunked_prefill == engine_args.get("enable_chunked_prefill") + assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling") legacy_parallel_config = engine_args.get("parallel_config") or {} assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get( "tensor_parallel_size", @@ -198,6 +204,64 @@ def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size +@pytest.mark.parametrize( + "cli_overrides", + [ + {"enable_lora": True}, + {"stage_0_enable_lora": True}, + ], + ids=["global", "stage-scoped"], +) +def test_from_pipeline_config_rejects_explicit_unowned_engine_cli_fields(cli_overrides): + with pytest.raises(ValueError, match=r"no structured config owner: enable_lora"): + _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides) + + +@pytest.mark.parametrize( + ("engine_extras", "unowned_field"), + [ + ({"enable_lora": True}, "enable_lora"), + ({"parallel_config": {"cfg_parallel_size": 2}}, "parallel_config.cfg_parallel_size"), + ], + ids=["top-level", "nested-parallel-config"], +) +def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras, unowned_field): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, engine_extras=engine_extras)]) + + with pytest.raises(ValueError, match=rf"no structured config owner: {unowned_field}"): + VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) + + +@pytest.mark.parametrize( + ("cli_overrides", "stage_id"), + [ + ({"kv_cache_dtype": "fp8"}, 0), + ({"stage_0_kv_cache_dtype": "fp8"}, 0), + ({"stage_1_kv_cache_dtype": "fp8"}, 1), + ], + ids=["global", "ar-stage", "generation-stage"], +) +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(): + stage = _from_pipeline_key( + "dreamzero", + deploy_config_path="dreamzero_tp1_cfg2", + cli_overrides={"kv_cache_dtype": "fp8"}, + ).stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.diffusion_kv_cache_dtype == "fp8" + + +def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged(): + assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True} + + def test_runtime_num_gpus_is_derived_from_parallel_world_size(): omni_config = _from_pipeline_key("hunyuan_image3_dit") stage = omni_config.stage_by_id(0) @@ -380,8 +444,10 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields -def test_runtime_config_fields_match_rfc_runtime_scope(): +def test_runtime_config_fields_match_structured_runtime_scope(): assert {f.name for f in fields(OmniStageRuntimeConfig)} == { + "distributed_executor_backend", + "worker_cls", "devices", "num_replicas", "env", @@ -392,8 +458,21 @@ def test_runtime_config_fields_match_rfc_runtime_scope(): } -def test_sub_config_fields_match_rfc_scopes(): +def test_sub_config_fields_match_structured_scopes(): assert {f.name for f in fields(OmniStageModelConfig)} == { + "model", + "model_arch", + "revision", + "tokenizer_revision", + "code_revision", + "seed", + "logits_processors", + "trust_remote_code", + "dtype", + "attention_backend", + "moe_backend", + "hf_overrides", + "limit_mm_per_prompt", "active_stream_window", "duplex_max_sessions", "enable_sleep_mode", @@ -404,6 +483,7 @@ def test_sub_config_fields_match_rfc_scopes(): "task_type", "codec_frame_rate_hz", "enforce_eager", + "max_cudagraph_capture_size", "enable_flashinfer_autotune", "compilation_config", "enable_multithread_weight_load", @@ -416,12 +496,16 @@ def test_sub_config_fields_match_rfc_scopes(): "tokenizer_subdir", } assert {f.name for f in fields(OmniStageLoadConfig)} == { + "tokenizer", + "download_dir", + "skip_tokenizer_init", "load_format", "tokenizer_mode", "config_format", "skip_mm_profiling", } assert {f.name for f in fields(OmniStageCacheConfig)} == { + "kv_cache_memory_bytes", "gpu_memory_utilization", "enable_prefix_caching", "disable_hybrid_kv_cache_manager", @@ -435,6 +519,8 @@ def test_sub_config_fields_match_rfc_scopes(): "async_scheduling", } assert {f.name for f in fields(OmniStageConnectorConfig)} == { + "async_chunk", + "omni_kv_config", "stage_connector", "output_connectors", "input_connectors", @@ -827,7 +913,7 @@ def test_from_pipeline_config_matches_build_engine_args_dict_behavior_for_repres pipeline = _resolve_pipeline_or_skip("qwen3_tts") legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] omega_stage = legacy_stage.to_omegaconf() - legacy_engine_args = build_engine_args_dict( + legacy_engine_args = build_legacy_engine_args_dict( omega_stage, model="/tmp/qwen3-tts", stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}}, @@ -922,6 +1008,8 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar def test_diffusion_config_field_classification_covers_current_fields(): + from vllm_omni.diffusion.data import OmniDiffusionConfig + classified_fields = ( omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS | omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS @@ -929,8 +1017,12 @@ def test_diffusion_config_field_classification_covers_current_fields(): ) assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)} + assert {f.name for f in fields(OmniDiffusionConfig)} <= ( + classified_fields | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS + ) assert { "enable_prompt_embed_cache", + "enable_session_state_manager", "prompt_embed_cache_size", "diffusion_kv_cache_dtype", } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index 2eb88bbd6..f81dee11c 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -3,15 +3,13 @@ import os import re import subprocess import threading -import time from datetime import datetime from pathlib import Path from typing import Any import pytest -from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw -from tests.helpers.mark import hardware_marks +from tests.helpers.mark import get_hardware_mark_list, hardware_marks from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import modify_stage_config from vllm_omni.platforms import current_omni_platform @@ -119,32 +117,6 @@ def create_unique_server_pytest_params( ] -def create_benchmark_pytest_params( - benchmark_configs: list[dict[str, Any]], - server_to_benchmark_mapping: dict[str, dict], -) -> list[Any]: - """Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks.""" - marks_by_name = _marks_by_test_name(benchmark_configs) - params: list[Any] = [] - seen: set[str] = set() - for config in benchmark_configs: - test_name = config["test_name"] - if test_name in seen: - continue - seen.add(test_name) - params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping) - id_suffixes = _unique_benchmark_param_id_suffixes(params_list) - for idx, id_suffix in enumerate(id_suffixes): - params.append( - pytest.param( - (test_name, idx), - marks=marks_by_name.get(test_name, []), - id=f"{test_name}-{id_suffix}", - ) - ) - return params - - def create_paired_benchmark_pytest_params( server_entries: list[tuple[Any, str]], params_by_test_name: dict[str, list[Any]], @@ -357,10 +329,6 @@ def supports_video_generation(model_name: str) -> bool: return any(key in lower for key in ("wan", "video", "i2v", "t2v")) -def supports_chat_generation(model_name: str) -> bool: - return not supports_video_generation(model_name) - - def parse_stage_devices(stage_config_path: str) -> str: text = Path(stage_config_path).read_text(encoding="utf-8") raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE) @@ -389,48 +357,6 @@ def assert_fault_exception(exc: Exception, error_keywords: tuple[str, ...]) -> N assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}" -def wait_chat_request_ready(host: str, port: int, model: str, timeout_sec: int = 180) -> None: - """Poll a minimal chat request until success.""" - deadline = time.time() + timeout_sec - payload = json.dumps( - { - "model": model, - "messages": [{"role": "user", "content": "Say hello in one short sentence."}], - "stream": False, - "modalities": ["text"], - } - ) - last_error: str | None = None - while time.time() < deadline: - try: - status, body = post_chat_completions_raw(host, port, payload) - if status == 200: - return - last_error = f"http={status} body={body[:200]!r}" - except Exception as exc: # noqa: BLE001 - last_error = str(exc) - time.sleep(2) - raise TimeoutError(f"runtime-teardown warmup request did not succeed within {timeout_sec}s: {last_error}") - - -def assert_no_extra_worker_processes( - baseline_pids: set[int], - worker_pattern: str, - timeout_sec: int = 60, -) -> None: - """Ensure no extra worker PIDs remain after teardown.""" - deadline = time.time() + timeout_sec - while time.time() < deadline: - current = set(list_remote_process_pids_by_pattern(worker_pattern)) - extra = current - baseline_pids - if not extra: - return - time.sleep(2) - current = set(list_remote_process_pids_by_pattern(worker_pattern)) - extra = sorted(current - baseline_pids) - assert not extra, f"orphan worker processes remain after container teardown: {extra}" - - def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]: mapping = {} for config in configs: @@ -520,36 +446,16 @@ def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> li return unique -def extract_mark_resource_label(mark_field: Any) -> str: - """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values. - - Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``. - - Prefer :func:`get_runtime_resource_label` for perf result filenames so labels - reflect the machine that actually ran the benchmark. - """ - if isinstance(mark_field, list): - for item in mark_field: - if isinstance(item, dict) and "hardware_marks" in item: - return extract_mark_resource_label(item) - return "na" - if not isinstance(mark_field, dict): - return "na" - hw = mark_field.get("hardware_marks") - if not isinstance(hw, dict): - return "na" - res = hw.get("res") - if not isinstance(res, dict) or not res: - return "na" - labels = [_safe_filename_token(value) for value in res.values()] - return "-".join(labels) if labels else "na" - - -_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = ( +# Device marketing-name tokens for ``get_runtime_resource_label`` only. +# Longer / more specific tokens must appear before shorter substrings +# (e.g. ``H200`` before ``H20``, ``910B4`` before ``910``). +_RUNTIME_DEVICE_ALIASES: tuple[str, ...] = ( "H100", "H800", "H200", "H20", + "B200", + "GB200", "L40S", "L40", "L4", @@ -578,7 +484,7 @@ def _normalize_runtime_device_label(raw: str) -> str: if not raw or not str(raw).strip(): return "na" upper = str(raw).upper() - for token in _KNOWN_RUNTIME_RESOURCE_TOKENS: + for token in _RUNTIME_DEVICE_ALIASES: if token.upper() in upper: return _safe_filename_token(token) compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw)) @@ -631,34 +537,41 @@ def resource_label_for_filename(resource_label: str | None) -> str: return token -def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str: - """Return runtime hardware label for perf result filenames.""" - del configs - return get_runtime_resource_label() +def is_hardware_nested_baseline(baseline: dict[str, Any]) -> bool: + """True for ``{"H100": {"metric": ...}, "A3": {...}}`` baselines. + + Detection is structural + hardware-label allowlist (not metric-name hints): + + - every top-level value is a non-empty dict (per-hardware metric map); + - every top-level key is in :func:`tests.helpers.mark.get_hardware_mark_list` + (``[hardware-resource]`` markers from ``pyproject.toml``). + + Metric names under each hardware bucket are unconstrained. + Runtime device aliases (:data:`_RUNTIME_DEVICE_ALIASES`) are unrelated. + """ + if not baseline: + return False + if not all(isinstance(value, dict) and bool(value) for value in baseline.values()): + return False + allowed = {label.upper() for label in get_hardware_mark_list()} + return all((token := str(key).strip()) and token.upper() in allowed for key in baseline.keys()) def resolve_baseline_value( baseline_raw: Any, *, sweep_index: int | None, - max_concurrency: Any = None, - request_rate: Any = None, ) -> Any: - """Pick the baseline threshold for this sweep step.""" + """Pick the baseline threshold for one metric at this sweep step. + + Accepts a sweep-aligned list/tuple (indexed by ``sweep_index``) or a scalar. + """ if baseline_raw is None: - return 100000 + return None if isinstance(baseline_raw, dict): - if max_concurrency is not None: - for key in (max_concurrency, str(max_concurrency)): - if key in baseline_raw: - return baseline_raw[key] - if request_rate is not None: - for key in (request_rate, str(request_rate)): - if key in baseline_raw: - return baseline_raw[key] - raise KeyError( - f"baseline dict has no key for max_concurrency={max_concurrency!r} " - f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}" + raise TypeError( + "per-metric baseline dict keyed by concurrency/request-rate is not supported; " + "use a sweep-aligned list or a scalar under each hardware bucket" ) if isinstance(baseline_raw, (list, tuple)): if sweep_index is None: @@ -669,22 +582,62 @@ def resolve_baseline_value( return baseline_raw -def _baseline_thresholds_for_step( - baseline_data: dict[str, Any], +def resolve_baseline_for_sweep( + baseline_config: dict[str, Any] | None, *, sweep_index: int | None = None, - max_concurrency: Any = None, - request_rate: Any = None, ) -> dict[str, Any]: - """Resolve baseline config to one threshold per metric for this iteration.""" - return { - metric_name: resolve_baseline_value( - baseline_raw, - sweep_index=sweep_index, - max_concurrency=max_concurrency, - request_rate=request_rate, + """Resolve sweep-dependent baselines while keeping every hardware bucket. + + Expects hardware-nested input only:: + + {"H100": {"throughput_qps": [a, b, c]}, "A3": {"throughput_qps": [x, y, z]}} + + with ``sweep_index=2`` (e.g. concurrency 32) becomes:: + + {"H100": {"throughput_qps": c}, "A3": {"throughput_qps": z}} + + Per-metric values may be a sweep-aligned list or a scalar. + """ + if not baseline_config: + return {} + if not is_hardware_nested_baseline(baseline_config): + allowed = sorted(get_hardware_mark_list()) + unknown = [ + str(key).strip() + for key in baseline_config.keys() + if str(key).strip().upper() not in {label.upper() for label in allowed} + ] + nested_shape = all(isinstance(value, dict) and bool(value) for value in baseline_config.values()) + if nested_shape and unknown: + raise ValueError( + "baseline top-level keys must be [hardware-resource] pytest markers from " + "pyproject.toml (tool.pytest.ini_options.markers). " + f"Unknown hardware label(s): {unknown!r}. " + f"Known markers: {allowed!r}. " + "Add a marker such as '\"A100: [hardware-resource] Tests that require A100 GPU\"' " + "to pyproject.toml, then retry." + ) + raise ValueError( + "baseline must be hardware-nested with [hardware-resource] marker names as " + "top-level keys, e.g. " + '{"H100": {"throughput_qps": [...]}, "A3": {"custom_metric": 1.0}}. ' + f"got top-level keys={list(baseline_config.keys())!r}; " + f"known markers={allowed!r}. " + "If you need a new hardware bucket, add it to pyproject.toml markers with " + "the [hardware-resource] tag." ) - for metric_name, baseline_raw in baseline_data.items() + + return { + hardware: { + metric_name: resolve_baseline_value( + baseline_raw, + sweep_index=sweep_index, + ) + for metric_name, baseline_raw in metrics.items() + } + for hardware, metrics in baseline_config.items() + if isinstance(metrics, dict) } @@ -697,17 +650,17 @@ def run_benchmark( *, baseline_config: dict[str, Any] | None = None, sweep_index: int | None = None, - request_rate: Any | None = None, - max_concurrency: Any | None = None, random_input_len: Any | None = None, random_output_len: Any | None = None, resource_label: str | None = None, ) -> dict[str, Any]: """Run one ``vllm bench serve --omni`` iteration and return parsed metrics. - After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds - (when ``baseline_config`` is provided). If the benchmark exits without writing a result file, - ``result_omni_template.json`` is used as a fallback. + After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the + sweep-resolved baseline (when ``baseline_config`` is provided). Hardware-nested + maps keep every hardware bucket, but each metric is reduced to the value for + this concurrency / request-rate step. If the benchmark exits without writing a + result file, ``result_omni_template.json`` is used as a fallback. """ current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") ri = _safe_filename_token(random_input_len) @@ -774,11 +727,10 @@ def run_benchmark( result = json.load(f) if baseline_config: - result["baseline"] = _baseline_thresholds_for_step( + # Keep every hardware bucket; resolve list metrics to this sweep step. + result["baseline"] = resolve_baseline_for_sweep( baseline_config, sweep_index=sweep_index, - request_rate=request_rate, - max_concurrency=max_concurrency, ) else: result["baseline"] = {} @@ -800,12 +752,3 @@ def pytest_addoption(parser: pytest.Parser) -> None: default=None, help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"), ) - parser.addoption( - "--assert-baseline", - action="store_true", - default=False, - help=( - "When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config " - "(default: off)." - ), - ) diff --git a/tests/dfx/perf/scripts/diffusion_result_template.json b/tests/dfx/perf/scripts/diffusion_result_template.json index ba16a993d..abb41c48b 100644 --- a/tests/dfx/perf/scripts/diffusion_result_template.json +++ b/tests/dfx/perf/scripts/diffusion_result_template.json @@ -21,10 +21,12 @@ "num-input-images": 0, "enable-negative-prompt": false, "baseline": { - "throughput_qps": 0, - "latency_mean": 0, - "peak_memory_mb_max": 0, - "peak_memory_mb_mean": 0 + "H100": { + "throughput_qps": 0, + "latency_mean": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0 + } } }, "result": { diff --git a/tests/dfx/perf/scripts/result_omni_template.json b/tests/dfx/perf/scripts/result_omni_template.json index 1d6132140..ca94be0f7 100644 --- a/tests/dfx/perf/scripts/result_omni_template.json +++ b/tests/dfx/perf/scripts/result_omni_template.json @@ -46,9 +46,11 @@ "median_audio_duration_s": 0, "p99_audio_duration_s": 0, "baseline": { - "mean_ttft_ms": 0, - "mean_audio_ttfp_ms": 0, - "mean_audio_rtf": 0 + "H100": { + "mean_ttft_ms": 0, + "mean_audio_ttfp_ms": 0, + "mean_audio_rtf": 0 + } }, "random_input_len": 0, "random_output_len": 0 diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index f37ac6ada..4ab7959d2 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -15,15 +15,10 @@ from tests.dfx.conftest import ( get_runtime_resource_label, is_diffusion_perf_config, load_benchmark_configs, - resolve_baseline_value, run_benchmark, ) from tests.helpers.runtime import OmniServer -# Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline`` -# (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are -# defined in the same module. -# # Optional JSON field ``mark`` is applied as pytest marks via # ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks": # {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``). @@ -162,36 +157,8 @@ def benchmark_params(request): } -def assert_result( - result, - params, - num_prompt, - *, - assert_baseline: bool, - sweep_index: int | None = None, - max_concurrency: Any | None = None, - request_rate: Any | None = None, -) -> None: +def assert_result(result, num_prompt) -> None: assert result["completed"] == num_prompt, "Request failures exist" - if not assert_baseline: - return - baseline_data = params.get("baseline", {}) - for metric_name, baseline_raw in baseline_data.items(): - current_value = result[metric_name] - baseline_value = resolve_baseline_value( - baseline_raw, - sweep_index=sweep_index, - max_concurrency=max_concurrency, - request_rate=request_rate, - ) - if "throughput" in metric_name: - if current_value <= baseline_value: - print( - f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}" - ) - else: - if current_value >= baseline_value: - print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}") @pytest.mark.benchmark @@ -200,7 +167,7 @@ def assert_result( paired_benchmark_params, indirect=["omni_server", "benchmark_params"], ) -def test_performance_benchmark(omni_server, benchmark_params, request): +def test_performance_benchmark(omni_server, benchmark_params): test_name = benchmark_params["test_name"] params = benchmark_params["params"] dataset_name = params.get("dataset_name", "") @@ -212,7 +179,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request): print(f"Running benchmark for model: {model}") print(f"Benchmark parameters: {benchmark_params}") - assert_baseline = request.config.getoption("--assert-baseline", default=False) resource_label = get_runtime_resource_label() def to_list(value, default=None): @@ -271,7 +237,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request): break # QPS / request-rate sweep - for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)): + for sweep_index, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)): args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)] result = run_benchmark( args=args, @@ -280,24 +246,15 @@ def test_performance_benchmark(omni_server, benchmark_params, request): dataset_name=dataset_name, num_prompt=num_prompt, baseline_config=params.get("baseline"), - sweep_index=i, - request_rate=qps, - max_concurrency=None, + sweep_index=sweep_index, random_input_len=params.get("random_input_len"), random_output_len=params.get("random_output_len"), resource_label=resource_label, ) - assert_result( - result, - params, - num_prompt, - assert_baseline=assert_baseline, - sweep_index=i, - request_rate=qps, - ) + assert_result(result, num_prompt) # concurrency test - for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)): + for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)): args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"] result = run_benchmark( args=args, @@ -306,18 +263,9 @@ def test_performance_benchmark(omni_server, benchmark_params, request): dataset_name=dataset_name, num_prompt=num_prompt, baseline_config=params.get("baseline"), - sweep_index=i, - request_rate=None, - max_concurrency=concurrency, + sweep_index=sweep_index, random_input_len=params.get("random_input_len"), random_output_len=params.get("random_output_len"), resource_label=resource_label, ) - assert_result( - result, - params, - num_prompt, - assert_baseline=assert_baseline, - sweep_index=i, - max_concurrency=concurrency, - ) + assert_result(result, num_prompt) diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 375d2a440..75d1bdc21 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -13,8 +13,6 @@ A config JSON file may be passed via --test-config-file. If omitted, every ``*.j pytest run_diffusion_benchmark.py -m "diffusion" pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off). - Optional JSON field ``mark`` is applied as pytest marks on that case via ``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``). @@ -46,6 +44,7 @@ from tests.dfx.conftest import ( get_runtime_resource_label, hardware_json_value, is_diffusion_perf_config, + resolve_baseline_for_sweep, resolve_pytest_marks, resource_label_for_filename, ) @@ -274,11 +273,6 @@ def _write_result_record(record: dict[str, Any]) -> Path: return target -def _append_to_aggregated_file(record: dict[str, Any]) -> None: - """Backward-compatible wrapper; prefer :func:`_write_result_record`.""" - _write_result_record(record) - - _server_lock = threading.Lock() # --------------------------------------------------------------------------- @@ -356,10 +350,49 @@ def _kill_process_tree(pid: int) -> None: def _resolve_offline_model(model: str) -> str: - """Under HF_HUB_OFFLINE, resolve a HF repo id to its local snapshot dir.""" + """Resolve a model string to a local directory. + + Local paths pass through unchanged. HF repo ids are resolved to their local + snapshot dir when ``HF_HUB_OFFLINE`` is set. ``repo_id/subfolder`` (e.g. + ``MiniMaxAI/MiniMax-H3/FL2VA``) is resolved to the matching partition + subfolder of the snapshot, downloading only that subfolder when online. + """ import huggingface_hub - if not model or os.path.isdir(model) or not huggingface_hub.constants.HF_HUB_OFFLINE: + if not model or os.path.isdir(model): + return model + + # MiniMax-H3 custom code uses relative imports across files, which breaks + # when the HF cache snapshot is symlinked (``get_class_from_dynamic_module`` + # resolves ``realpath`` into the blobs dir). Mirror the accuracy test's env + # overrides so CI / local runs can point at a materialized partition. The + # repo-root key (used with ``--task-type fl2va|ref2va``) expects a + # materialized root containing the FL2VA/Ref2VA subfolders. + model_env_overrides = { + "MiniMaxAI/MiniMax-H3": "VLLM_TEST_MINIMAX_H3_MODEL", + "MiniMaxAI/MiniMax-H3/FL2VA": "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL", + "MiniMaxAI/MiniMax-H3/Ref2VA": "VLLM_TEST_MINIMAX_H3_REF2VA_MODEL", + } + env_name = model_env_overrides.get(model) + if env_name: + env_model = os.environ.get(env_name) + if env_model: + return env_model + + parts = model.split("/") + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + subfolder = "/".join(parts[2:]) + from huggingface_hub import snapshot_download + + snapshot_root = snapshot_download( + repo_id, + allow_patterns=[f"{subfolder}/**"], + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + return str(Path(snapshot_root) / subfolder) + + if not huggingface_hub.constants.HF_HUB_OFFLINE: return model from huggingface_hub import snapshot_download @@ -393,6 +426,11 @@ class DiffusionServer: def _start_server(self) -> None: env = os.environ.copy() env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + # Give each run its own writable video-storage dir. The server default + # (/tmp/storage) may be missing or not writable on shared hosts / CI. + storage_path = Path(os.environ.get("VLLM_OMNI_STORAGE_PATH", str(BENCHMARK_RESULT_DIR / "storage"))) + storage_path.mkdir(parents=True, exist_ok=True) + env["VLLM_OMNI_STORAGE_PATH"] = str(storage_path) cmd = [ sys.executable, @@ -678,7 +716,7 @@ def run_benchmark( with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp: tmp_result_file = Path(tmp.name) - exclude_keys = {"baseline", "dataset", "task", "name", "skip-performance-assertion"} + exclude_keys = {"baseline", "dataset", "task", "name"} cmd = [ sys.executable, @@ -772,12 +810,17 @@ def run_benchmark( completed = metrics.get("completed_requests", metrics.get("completed", 0)) failed = metrics.get("failed_requests", metrics.get("failed", 0)) + # Persist sweep-resolved baseline from params (already narrowed in _build_run_params). + baseline = params.get("baseline") or {} + metrics["baseline"] = baseline + record: dict[str, Any] = { "test_name": test_name, "endpoint": endpoint, "timestamp": timestamp, "server_params": server_cfg.get("server_params"), "benchmark_params": params, + "baseline": baseline, "result": metrics, "log_file": str(log_file), "Model": model, @@ -827,70 +870,14 @@ def run_benchmark( # --------------------------------------------------------------------------- -def _resolve_baseline_value( - baseline_raw: Any, - *, - sweep_index: int | None, - max_concurrency: Any = None, - request_rate: Any = None, -) -> Any: - """Pick the baseline threshold for this sweep step.""" - if baseline_raw is None: - return 100000 - if isinstance(baseline_raw, dict): - if max_concurrency is not None: - for key in (max_concurrency, str(max_concurrency)): - if key in baseline_raw: - return baseline_raw[key] - if request_rate is not None: - for key in (request_rate, str(request_rate)): - if key in baseline_raw: - return baseline_raw[key] - raise KeyError( - f"baseline dict has no key for max_concurrency={max_concurrency!r} " - f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}" - ) - if isinstance(baseline_raw, (list, tuple)): - if sweep_index is None: - raise ValueError("sweep_index is required when baseline is a list or tuple") - return baseline_raw[sweep_index] - return baseline_raw - - def assert_result( result: dict[str, Any], - params: dict[str, Any], num_prompts: int, - *, - sweep_index: int | None = None, - max_concurrency: Any = None, - request_rate: Any = None, - assert_baseline: bool = True, ) -> None: - """Assert that benchmark metrics satisfy the configured baselines.""" + """Assert that the expected number of requests completed.""" completed = result.get("completed_requests", result.get("completed", 0)) assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}" - if not assert_baseline: - return - if params.get("skip-performance-assertion", False): - print("Skipping performance assertions.") - return - - for metric, baseline_raw in params.get("baseline", {}).items(): - current = result.get(metric) - assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}" - threshold = _resolve_baseline_value( - baseline_raw, - sweep_index=sweep_index, - max_concurrency=max_concurrency, - request_rate=request_rate, - ) - if "throughput" in metric: - assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}" - else: - assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}" - def _default_benchmark_endpoint_for_task(task: str) -> str: """Return the default client-side benchmark endpoint for a diffusion task.""" @@ -932,15 +919,11 @@ def _build_run_params( if max_concurrency is not None: run_params["max-concurrency"] = max_concurrency if "baseline" in params: - run_params["baseline"] = { - metric: _resolve_baseline_value( - baseline_raw, - sweep_index=sweep_index, - max_concurrency=max_concurrency, - request_rate=request_rate, - ) - for metric, baseline_raw in params["baseline"].items() - } + # Keep all hardware buckets; pick the metric value for this sweep step only. + run_params["baseline"] = resolve_baseline_for_sweep( + params.get("baseline"), + sweep_index=sweep_index, + ) return run_params @@ -963,36 +946,30 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]: sweep_runs: list[dict[str, Any]] = [] - for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)): + for sweep_index, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)): sweep_runs.append( { "params": _build_run_params( params, request_rate=request_rate, num_prompts=num_prompts, - sweep_index=i, + sweep_index=sweep_index, ), "num_prompts": num_prompts, - "sweep_index": i, - "request_rate": request_rate, - "max_concurrency": None, } ) - for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)): + for sweep_index, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)): sweep_runs.append( { "params": _build_run_params( params, max_concurrency=max_concurrency, num_prompts=num_prompts, - sweep_index=i, request_rate="inf", + sweep_index=sweep_index, ), "num_prompts": num_prompts, - "sweep_index": i, - "request_rate": None, - "max_concurrency": max_concurrency, } ) @@ -1003,11 +980,9 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]: "params": _build_run_params( params, num_prompts=default_num_prompts, + sweep_index=0, ), "num_prompts": default_num_prompts, - "sweep_index": None, - "request_rate": None, - "max_concurrency": None, } ) @@ -1023,14 +998,12 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]: paired_benchmark_params, indirect=["diffusion_server", "benchmark_params"], ) -def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request): +def test_diffusion_performance_benchmark(diffusion_server, benchmark_params): """Run the diffusion performance benchmark and verify request completion. One server is started per unique parallel configuration (module scope). For each server, all benchmark parameter sets defined in the config JSON are executed sequentially; metrics are recorded to the aggregated result file. - - Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block. """ test_name = benchmark_params["test_name"] params = benchmark_params["params"] @@ -1073,14 +1046,4 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req print(f"\n Aggregated results: {_aggregated_result_file_for_source(str(source))}") print("=" * 60) - assert_baseline = request.config.getoption("--assert-baseline", default=False) - - assert_result( - result, - params, - sweep_run["num_prompts"], - sweep_index=sweep_run["sweep_index"], - max_concurrency=sweep_run["max_concurrency"], - request_rate=sweep_run["request_rate"], - assert_baseline=assert_baseline, - ) + assert_result(result, sweep_run["num_prompts"]) diff --git a/tests/dfx/perf/tests/test_bagel_vllm_omni.json b/tests/dfx/perf/tests/test_bagel_vllm_omni.json index e6777acfe..9f6ff0241 100644 --- a/tests/dfx/perf/tests/test_bagel_vllm_omni.json +++ b/tests/dfx/perf/tests/test_bagel_vllm_omni.json @@ -40,10 +40,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.495, - "latency_mean": 2.0238, - "peak_memory_mb_max": 29393.4286, - "peak_memory_mb_mean": 29392.7143 + "H100": { + "throughput_qps": 0.495, + "latency_mean": 2.0238, + "peak_memory_mb_max": 29393.4286, + "peak_memory_mb_mean": 29392.7143 + } } } ] @@ -90,10 +92,12 @@ "enable-negative-prompt": true, "num-input-images": 1, "baseline": { - "throughput_qps": 0.3451, - "latency_mean": 2.9031, - "peak_memory_mb_max": 30480.2857, - "peak_memory_mb_mean": 30479.2857 + "H100": { + "throughput_qps": 0.3451, + "latency_mean": 2.9031, + "peak_memory_mb_max": 30480.2857, + "peak_memory_mb_mean": 30479.2857 + } } } ] @@ -139,10 +143,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.5048, - "latency_mean": 1.987, - "peak_memory_mb_max": 29380.0, - "peak_memory_mb_mean": 29380.0 + "H100": { + "throughput_qps": 0.5048, + "latency_mean": 1.987, + "peak_memory_mb_max": 29380.0, + "peak_memory_mb_mean": 29380.0 + } } } ] @@ -189,10 +195,12 @@ "enable-negative-prompt": true, "num-input-images": 1, "baseline": { - "throughput_qps": 0.149, - "latency_mean": 6.7157, - "peak_memory_mb_max": 30574.0, - "peak_memory_mb_mean": 30574.0 + "H100": { + "throughput_qps": 0.149, + "latency_mean": 6.7157, + "peak_memory_mb_max": 30574.0, + "peak_memory_mb_mean": 30574.0 + } } } ] diff --git a/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json index 028889af1..e2f96f626 100644 --- a/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json +++ b/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json @@ -24,10 +24,12 @@ "warmup-num-inference-steps": 28, "enable-negative-prompt": true, "baseline": { - "throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302}, - "latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7}, - "peak_memory_mb_max": 40500, - "peak_memory_mb_mean": 40500 + "H100": { + "throughput_qps": [0.0303, 0.0302, 0.0302], + "latency_mean": [32.6, 62.2, 111.7], + "peak_memory_mb_max": 40500, + "peak_memory_mb_mean": 40500 + } } } ] diff --git a/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json b/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json index 4ee6ca577..22218268e 100644 --- a/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json @@ -24,10 +24,12 @@ "warmup-num-inference-steps": 28, "enable-negative-prompt": true, "baseline": { - "throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549}, - "latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4}, - "peak_memory_mb_max": 40500, - "peak_memory_mb_mean": 40500 + "H100": { + "throughput_qps": [0.0549, 0.0547, 0.0549], + "latency_mean": [18.0, 34.3, 61.4], + "peak_memory_mb_max": 40500, + "peak_memory_mb_mean": 40500 + } } } ] diff --git a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json index 9bf31658d..c4d2305e6 100644 --- a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json +++ b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json @@ -48,9 +48,11 @@ "guardrails": false }, "baseline": { - "throughput_qps": 1.107, - "latency_mean": 0.8929, - "peak_memory_mb_mean": 80000.0 + "H100": { + "throughput_qps": 1.107, + "latency_mean": 0.8929, + "peak_memory_mb_mean": 80000.0 + } } } ] @@ -115,9 +117,11 @@ "use_duration_template": false }, "baseline": { - "throughput_qps": 0.027, - "latency_mean": 36.6824, - "peak_memory_mb_mean": 30155.4 + "H100": { + "throughput_qps": 0.027, + "latency_mean": 36.6824, + "peak_memory_mb_mean": 30155.4 + } } } ] @@ -181,9 +185,11 @@ "guardrails": false }, "baseline": { - "throughput_qps": 0.027, - "latency_mean": 36.968, - "peak_memory_mb_mean": 30443.6 + "H100": { + "throughput_qps": 0.027, + "latency_mean": 36.968, + "peak_memory_mb_mean": 30443.6 + } } } ] @@ -251,9 +257,11 @@ "condition_video_keep": "first" }, "baseline": { - "throughput_qps": 0.027, - "latency_mean": 37.2117, - "peak_memory_mb_mean": 30135.6 + "H100": { + "throughput_qps": 0.027, + "latency_mean": 37.2117, + "peak_memory_mb_mean": 30135.6 + } } } ] diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json index 04cc88ac2..75e5983b0 100644 --- a/tests/dfx/perf/tests/test_higgs_audio_v3.json +++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json @@ -46,12 +46,14 @@ "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 223.3584 - ], - "median_audio_rtf": [ - 0.4393 - ] + "H100": { + "median_audio_ttfp_ms": [ + 223.3584 + ], + "median_audio_rtf": [ + 0.4393 + ] + } } }, { @@ -70,15 +72,17 @@ "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 276.575 - ], - "median_audio_rtf": [ - 0.5073 - ], - "audio_throughput": [ - 27.3728 - ] + "H100": { + "median_audio_ttfp_ms": [ + 276.575 + ], + "median_audio_rtf": [ + 0.5073 + ], + "audio_throughput": [ + 27.3728 + ] + } } } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json index d2a20af7a..c4ef34b96 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json +++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json @@ -133,13 +133,14 @@ "bot-task": "think_recaption", "max-concurrency": 16, "num-prompts": 96, - "skip-performance-assertion": true, "baseline": { - "throughput_qps": 0.1874, - "latency_p99": 101.852, - "latency_mean": 79.0385, - "peak_memory_mb_max": 70000, - "peak_memory_mb_mean": 70000 + "H200": { + "throughput_qps": 0.1874, + "latency_p99": 101.852, + "latency_mean": 79.0385, + "peak_memory_mb_max": 70000, + "peak_memory_mb_mean": 70000 + } } } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json index 733579e96..a86a6a3d5 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json @@ -37,12 +37,13 @@ "num-inference-steps": 8, "num-prompts": 10, "max-concurrency": 1, - "skip-performance-assertion": true, "baseline": { - "throughput_qps": 0.21, - "latency_p99": 4.7469, - "peak_memory_mb_max": 101100, - "peak_memory_mb_mean": 101100 + "H200": { + "throughput_qps": 0.21, + "latency_p99": 4.7469, + "peak_memory_mb_max": 101100, + "peak_memory_mb_mean": 101100 + } } } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json index 27b04a469..e49cd0543 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json @@ -37,12 +37,13 @@ "num-inference-steps": 8, "num-prompts": 10, "max-concurrency": 1, - "skip-performance-assertion": true, "baseline": { - "throughput_qps": 0.20, - "latency_p99": 5.1025, - "peak_memory_mb_max": 97402, - "peak_memory_mb_mean": 97402 + "H200": { + "throughput_qps": 0.20, + "latency_p99": 5.1025, + "peak_memory_mb_max": 97402, + "peak_memory_mb_mean": 97402 + } } } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json index 28842164c..54631471a 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json @@ -36,13 +36,14 @@ "num-inference-steps": 8, "num-prompts": 10, "max-concurrency": 1, - "skip-performance-assertion": true, "baseline": { - "throughput_qps": 0.213, - "latency_mean": 4.6953, - "latency_p99": 4.8601, - "peak_memory_mb_max": 56912.0, - "peak_memory_mb_mean": 56912.0 + "H200": { + "throughput_qps": 0.213, + "latency_mean": 4.6953, + "latency_p99": 4.8601, + "peak_memory_mb_max": 56912.0, + "peak_memory_mb_mean": 56912.0 + } } } ] diff --git a/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json b/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json index a6751bda1..07403e78f 100644 --- a/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json +++ b/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json @@ -26,8 +26,7 @@ "warmup-requests": 1, "warmup-concurrency": 1, "warmup-num-inference-steps": 2, - "enable-negative-prompt": true, - "skip-performance-assertion": true + "enable-negative-prompt": true } ] } diff --git a/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json b/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json new file mode 100644 index 000000000..a28aab96c --- /dev/null +++ b/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json @@ -0,0 +1,295 @@ +[ + { + "test_name": "test_minimax_h3_t2v_4gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "MiniMax-H3 T2VA text-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "MiniMaxAI/MiniMax-H3", + "serve_args": { + "trust-remote-code": true, + "task-type": "fl2va", + "num-gpus": 4, + "usp": 4, + "ring": 1, + "use-hsdp": true, + "hsdp-shard-size": 4, + "text-encoder-tp-size": 4, + "vae-patch-parallel-size": 4, + "vae-parallel-mode": "tile", + "vae-use-tiling": true, + "diffusion-attention-backend": "FLASH_ATTN", + "stage-init-timeout": 1800, + "init-timeout": 1800, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1344x768_frames209_steps8", + "dataset": "random", + "task": "t2v", + "num-prompts": 3, + "max-concurrency": 1, + "warmup-requests": 1, + "warmup-concurrency": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1344, + "height": 768, + "prompt": "Cinematic drone shot pushing through misty mountain ridges at sunrise, golden light catching the clouds, slow and smooth.", + "num_inference_steps": 8, + "num_frames": 209, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "flow_shift": 12.0, + "extra_params": "{\"task\":\"t2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}" + }, + "baseline": { + "H100": { + "throughput_qps": 0.0211, + "latency_mean": 38.3252, + "peak_memory_mb_mean": 54219.6 + } + } + } + ] + }, + { + "test_name": "test_minimax_h3_ti2v_4gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "MiniMax-H3 FL2VA first-frame image-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "MiniMaxAI/MiniMax-H3", + "serve_args": { + "trust-remote-code": true, + "task-type": "fl2va", + "num-gpus": 4, + "usp": 4, + "ring": 1, + "use-hsdp": true, + "hsdp-shard-size": 4, + "text-encoder-tp-size": 4, + "vae-patch-parallel-size": 4, + "vae-parallel-mode": "tile", + "vae-use-tiling": true, + "diffusion-attention-backend": "FLASH_ATTN", + "stage-init-timeout": 1800, + "init-timeout": 1800, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1344x768_frames209_steps8", + "dataset": "random", + "task": "ti2v", + "num-prompts": 3, + "max-concurrency": 1, + "warmup-requests": 1, + "warmup-concurrency": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1344, + "height": 768, + "prompt": "Continue the scene: the camera slowly pushes in on the subject as the background city lights drift and blur, with the actor singing softly and emotionally.", + "num_inference_steps": 8, + "num_frames": 209, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "flow_shift": 12.0, + "extra_params": "{\"task\":\"fl2va\",\"duration\":8.7,\"audio_flow_shift\":3.0}" + }, + "num-input-images": 1, + "baseline": { + "H100": { + "throughput_qps": 0.0212, + "latency_mean": 38.1385, + "peak_memory_mb_mean": 54622.8 + } + } + } + ] + }, + { + "test_name": "test_minimax_h3_v2v_4gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "MiniMax-H3 Ref2VA video-reference-to-video+audio on 4 GPUs (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "MiniMaxAI/MiniMax-H3", + "serve_args": { + "trust-remote-code": true, + "task-type": "ref2va", + "num-gpus": 4, + "usp": 4, + "ring": 1, + "use-hsdp": true, + "hsdp-shard-size": 4, + "text-encoder-tp-size": 4, + "vae-patch-parallel-size": 4, + "vae-parallel-mode": "tile", + "vae-use-tiling": true, + "diffusion-attention-backend": "FLASH_ATTN", + "stage-init-timeout": 1800, + "init-timeout": 1800, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1344x768_frames209_steps8", + "dataset": "random", + "task": "v2v", + "num-prompts": 3, + "max-concurrency": 1, + "warmup-requests": 1, + "warmup-concurrency": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1344, + "height": 768, + "prompt": "Restyle the reference footage into a cinematic night scene with neon reflections, keeping the same camera motion and subject.", + "num_inference_steps": 8, + "num_frames": 209, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "flow_shift": 12.0, + "extra_params": "{\"task\":\"ref2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}" + }, + "baseline": { + "H100": { + "throughput_qps": 0.0077, + "latency_mean": 105.7508, + "peak_memory_mb_mean": 55180.8 + } + } + } + ] + }, + { + "test_name": "test_minimax_h3_t2v_dlo_4gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "MiniMax-H3 T2VA with 4-card distributed layerwise offload (DLO), 8 denoise steps (USP4 + HSDP4, encoder-TP4, VAE-PP4 tile)", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "MiniMaxAI/MiniMax-H3", + "serve_args": { + "trust-remote-code": true, + "task-type": "fl2va", + "num-gpus": 4, + "usp": 4, + "ring": 1, + "use-hsdp": true, + "hsdp-shard-size": 4, + "text-encoder-tp-size": 4, + "vae-patch-parallel-size": 4, + "vae-parallel-mode": "tile", + "vae-use-tiling": true, + "diffusion-attention-backend": "FLASH_ATTN", + "stage-init-timeout": 1800, + "init-timeout": 1800, + "enable-diffusion-pipeline-profiler": true, + "enable-distributed-layerwise-offload": true, + "dlo-no-use-allgather": true, + "dlo-resident-layers": 8, + "enforce-eager": true + } + }, + "benchmark_params": [ + { + "name": "1344x768_frames209_steps8", + "dataset": "random", + "task": "t2v", + "num-prompts": 3, + "max-concurrency": 1, + "warmup-requests": 1, + "warmup-concurrency": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1344, + "height": 768, + "prompt": "Cinematic drone shot pushing through misty mountain ridges at sunrise, golden light catching the clouds, slow and smooth.", + "num_inference_steps": 8, + "num_frames": 209, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "flow_shift": 12.0, + "extra_params": "{\"task\":\"t2va\",\"duration\":8.7,\"aspect_ratio\":\"16:9\",\"audio_flow_shift\":3.0}" + }, + "baseline": { + "H100": { + "throughput_qps": 0.0212, + "latency_mean": 38.2089, + "peak_memory_mb_mean": 36218.4 + } + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json index be108cbd6..a73abc5fb 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json @@ -44,27 +44,29 @@ "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 96.4061, - 140.8147, - 271.9025, - 362.2916, - 507.776 - ], - "mean_e2el_ms": [ - 18507.4582, - 28365.4956, - 31907.9613, - 48161.5829, - 72630.8601 - ], - "mean_audio_rtf": [ - 0.1175, - 0.18, - 0.2359, - 0.3278, - 0.6296 - ] + "H100": { + "mean_ttft_ms": [ + 96.4061, + 140.8147, + 271.9025, + 362.2916, + 507.776 + ], + "mean_e2el_ms": [ + 18507.4582, + 28365.4956, + 31907.9613, + 48161.5829, + 72630.8601 + ], + "mean_audio_rtf": [ + 0.1175, + 0.18, + 0.2359, + 0.3278, + 0.6296 + ] + } } }, { @@ -91,15 +93,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 354.0951 - ], - "mean_e2el_ms": [ - 3952.5615 - ], - "mean_audio_rtf": [ - 0.117 - ] + "H100": { + "mean_ttft_ms": [ + 354.0951 + ], + "mean_e2el_ms": [ + 3952.5615 + ], + "mean_audio_rtf": [ + 0.117 + ] + } } }, { @@ -128,15 +132,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 259.6317 - ], - "mean_e2el_ms": [ - 4253.3009 - ], - "mean_audio_rtf": [ - 0.1448 - ] + "H100": { + "mean_ttft_ms": [ + 259.6317 + ], + "mean_e2el_ms": [ + 4253.3009 + ], + "mean_audio_rtf": [ + 0.1448 + ] + } } }, { @@ -167,15 +173,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 457.4453 - ], - "mean_e2el_ms": [ - 5912.5623 - ], - "mean_audio_rtf": [ - 0.1967 - ] + "H100": { + "mean_ttft_ms": [ + 457.4453 + ], + "mean_e2el_ms": [ + 5912.5623 + ], + "mean_audio_rtf": [ + 0.1967 + ] + } } }, { @@ -206,20 +214,22 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 90.8033, - 202.1779, - 326.5864, - 544.8914, - 961.3929 - ], - "mean_e2el_ms": [ - 4517.3411, - 7857.9711, - 9861.1687, - 13283.3363, - 19541.0736 - ] + "H100": { + "mean_ttft_ms": [ + 90.8033, + 202.1779, + 326.5864, + 544.8914, + 961.3929 + ], + "mean_e2el_ms": [ + 4517.3411, + 7857.9711, + 9861.1687, + 13283.3363, + 19541.0736 + ] + } } }, { @@ -251,12 +261,14 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 318.048 - ], - "mean_e2el_ms": [ - 807.6114 - ] + "H100": { + "mean_ttft_ms": [ + 318.048 + ], + "mean_e2el_ms": [ + 807.6114 + ] + } } }, { @@ -290,12 +302,14 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 186.168 - ], - "mean_e2el_ms": [ - 726.1137 - ] + "H100": { + "mean_ttft_ms": [ + 186.168 + ], + "mean_e2el_ms": [ + 726.1137 + ] + } } }, { @@ -331,12 +345,14 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 375.5282 - ], - "mean_e2el_ms": [ - 1006.1663 - ] + "H100": { + "mean_ttft_ms": [ + 375.5282 + ], + "mean_e2el_ms": [ + 1006.1663 + ] + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json index 4e54c6f16..a3d8aaea8 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json @@ -35,10 +35,12 @@ "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085], - "mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485], - "mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275], - "mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322] + "H100": { + "mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085], + "mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485], + "mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275], + "mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322] + } } }, { @@ -61,9 +63,11 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [340.5122], - "mean_audio_ttfp_ms": [463.0752], - "mean_audio_rtf": [0.1153] + "H100": { + "mean_ttft_ms": [340.5122], + "mean_audio_ttfp_ms": [463.0752], + "mean_audio_rtf": [0.1153] + } } }, { @@ -88,9 +92,11 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [236.926], - "mean_audio_ttfp_ms": [395.2182], - "mean_audio_rtf": [0.1259] + "H100": { + "mean_ttft_ms": [236.926], + "mean_audio_ttfp_ms": [395.2182], + "mean_audio_rtf": [0.1259] + } } }, { @@ -117,9 +123,11 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [470.6126], - "mean_audio_ttfp_ms": [665.5474], - "mean_audio_rtf": [0.166] + "H100": { + "mean_ttft_ms": [470.6126], + "mean_audio_ttfp_ms": [665.5474], + "mean_audio_rtf": [0.166] + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json index 3fbeafb9d..38dbfb2e3 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json @@ -44,27 +44,29 @@ "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 115.8348, - 173.1823, - 446.672, - 464.5985, - 3420.0634 - ], - "mean_e2el_ms": [ - 25070.9523, - 31992.0646, - 37619.7669, - 50393.8257, - 136619.2743 - ], - "mean_audio_rtf": [ - 0.169, - 0.2111, - 0.2475, - 0.3551, - 0.9266 - ] + "H100": { + "mean_ttft_ms": [ + 115.8348, + 173.1823, + 446.672, + 464.5985, + 3420.0634 + ], + "mean_e2el_ms": [ + 25070.9523, + 31992.0646, + 37619.7669, + 50393.8257, + 136619.2743 + ], + "mean_audio_rtf": [ + 0.169, + 0.2111, + 0.2475, + 0.3551, + 0.9266 + ] + } } }, { @@ -91,15 +93,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 343.1377 - ], - "mean_e2el_ms": [ - 4811.1123 - ], - "mean_audio_rtf": [ - 0.1463 - ] + "H100": { + "mean_ttft_ms": [ + 343.1377 + ], + "mean_e2el_ms": [ + 4811.1123 + ], + "mean_audio_rtf": [ + 0.1463 + ] + } } }, { @@ -128,15 +132,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 245.5446 - ], - "mean_e2el_ms": [ - 5019.1256 - ], - "mean_audio_rtf": [ - 0.1713 - ] + "H100": { + "mean_ttft_ms": [ + 245.5446 + ], + "mean_e2el_ms": [ + 5019.1256 + ], + "mean_audio_rtf": [ + 0.1713 + ] + } } }, { @@ -167,15 +173,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [ - 506.1407 - ], - "mean_e2el_ms": [ - 6814.9656 - ], - "mean_audio_rtf": [ - 0.2213 - ] + "H100": { + "mean_ttft_ms": [ + 506.1407 + ], + "mean_e2el_ms": [ + 6814.9656 + ], + "mean_audio_rtf": [ + 0.2213 + ] + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json index 4dedcd4b9..8775ee381 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json @@ -51,20 +51,22 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 69.8081, - 173.0651, - 271.8108, - 413.4079, - 497.2546 - ], - "mean_e2el_ms": [ - 4720.8225, - 7937.2532, - 9962.3566, - 13724.5479, - 18324.091 - ] + "H100": { + "mean_ttft_ms": [ + 69.8081, + 173.0651, + 271.8108, + 413.4079, + 497.2546 + ], + "mean_e2el_ms": [ + 4720.8225, + 7937.2532, + 9962.3566, + 13724.5479, + 18324.091 + ] + } } }, { @@ -97,12 +99,14 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 25.9477 - ], - "mean_e2el_ms": [ - 517.3155 - ] + "H100": { + "mean_ttft_ms": [ + 25.9477 + ], + "mean_e2el_ms": [ + 517.3155 + ] + } } }, { @@ -137,12 +141,14 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 26.7056 - ], - "mean_e2el_ms": [ - 551.308 - ] + "H100": { + "mean_ttft_ms": [ + 26.7056 + ], + "mean_e2el_ms": [ + 551.308 + ] + } } }, { @@ -179,12 +185,14 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [ - 27.6906 - ], - "mean_e2el_ms": [ - 596.5284 - ] + "H100": { + "mean_ttft_ms": [ + 27.6906 + ], + "mean_e2el_ms": [ + 596.5284 + ] + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json index 75a3d7744..dcd2188c5 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json @@ -34,10 +34,12 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0721, - "latency_mean": 13.8398, - "peak_memory_mb_max": 57610.0, - "peak_memory_mb_mean": 57610.0 + "H100": { + "throughput_qps": 0.0721, + "latency_mean": 13.8398, + "peak_memory_mb_max": 57610.0, + "peak_memory_mb_mean": 57610.0 + } } }, { @@ -52,10 +54,12 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0186, - "latency_mean": 53.6942, - "peak_memory_mb_max": 67790.0, - "peak_memory_mb_mean": 67790.0 + "H100": { + "throughput_qps": 0.0186, + "latency_mean": 53.6942, + "peak_memory_mb_max": 67790.0, + "peak_memory_mb_mean": 67790.0 + } } } ] @@ -100,10 +104,12 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0538, - "latency_mean": 18.5591, - "peak_memory_mb_max": 56449.4286, - "peak_memory_mb_mean": 56449.4286 + "H100": { + "throughput_qps": 0.0538, + "latency_mean": 18.5591, + "peak_memory_mb_max": 56449.4286, + "peak_memory_mb_mean": 56449.4286 + } } } ] @@ -157,10 +163,12 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.2776, - "latency_mean": 3.5938, - "peak_memory_mb_max": 58113.7143, - "peak_memory_mb_mean": 58075.1429 + "H100": { + "throughput_qps": 0.2776, + "latency_mean": 3.5938, + "peak_memory_mb_max": 58113.7143, + "peak_memory_mb_mean": 58075.1429 + } } }, { @@ -175,10 +183,12 @@ "num-input-images": 2, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1051, - "latency_mean": 9.4918, - "peak_memory_mb_max": 68107.4286, - "peak_memory_mb_mean": 68107.4286 + "H100": { + "throughput_qps": 0.1051, + "latency_mean": 9.4918, + "peak_memory_mb_max": 68107.4286, + "peak_memory_mb_mean": 68107.4286 + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json index 9c5ad947b..686bd6160 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json @@ -33,10 +33,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1129, - "latency_mean": 8.8352, - "peak_memory_mb_max": 58572.0, - "peak_memory_mb_mean": 58572.0 + "H100": { + "throughput_qps": 0.1129, + "latency_mean": 8.8352, + "peak_memory_mb_max": 58572.0, + "peak_memory_mb_mean": 58572.0 + } } }, { @@ -50,10 +52,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.024, - "latency_mean": 41.5686, - "peak_memory_mb_max": 66368.0, - "peak_memory_mb_mean": 66368.0 + "H100": { + "throughput_qps": 0.024, + "latency_mean": 41.5686, + "peak_memory_mb_max": 66368.0, + "peak_memory_mb_mean": 66368.0 + } } } ] @@ -96,10 +100,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0672, - "latency_mean": 14.8517, - "peak_memory_mb_max": 57076.0, - "peak_memory_mb_mean": 57076.0 + "H100": { + "throughput_qps": 0.0672, + "latency_mean": 14.8517, + "peak_memory_mb_max": 57076.0, + "peak_memory_mb_mean": 57076.0 + } } } ] @@ -152,10 +158,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.2879, - "latency_mean": 3.4701, - "peak_memory_mb_max": 58758.0, - "peak_memory_mb_mean": 58758.0 + "H100": { + "throughput_qps": 0.2879, + "latency_mean": 3.4701, + "peak_memory_mb_max": 58758.0, + "peak_memory_mb_mean": 58758.0 + } } }, { @@ -169,10 +177,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1113, - "latency_mean": 8.9656, - "peak_memory_mb_max": 67424.0, - "peak_memory_mb_mean": 67424.0 + "H100": { + "throughput_qps": 0.1113, + "latency_mean": 8.9656, + "peak_memory_mb_max": 67424.0, + "peak_memory_mb_mean": 67424.0 + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json index 0553d891e..bc4d93a24 100644 --- a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -33,10 +33,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0703, - "latency_mean": 14.208, - "peak_memory_mb_max": 63225.1429, - "peak_memory_mb_mean": 63225.1429 + "H100": { + "throughput_qps": 0.0703, + "latency_mean": 14.208, + "peak_memory_mb_max": 63225.1429, + "peak_memory_mb_mean": 63225.1429 + } } }, { @@ -50,10 +52,12 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0417, - "latency_mean": 23.914, - "peak_memory_mb_max": 63225.1429, - "peak_memory_mb_mean": 63225.1429 + "H100": { + "throughput_qps": 0.0417, + "latency_mean": 23.914, + "peak_memory_mb_max": 63225.1429, + "peak_memory_mb_mean": 63225.1429 + } } } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index 3fb882563..ebf9ccc29 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -33,9 +33,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.4153, - "latency_mean": 2.4065, - "peak_memory_mb_mean": 55121.4286 + "H100": { + "throughput_qps": 0.4153, + "latency_mean": 2.4065, + "peak_memory_mb_mean": 55121.4286 + } } }, { @@ -49,9 +51,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0423, - "latency_mean": 23.5991, - "peak_memory_mb_mean": 65118.0 + "H100": { + "throughput_qps": 0.0423, + "latency_mean": 23.5991, + "peak_memory_mb_mean": 65118.0 + } } } ] @@ -91,9 +95,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.4175, - "latency_mean": 2.3953, - "peak_memory_mb_mean": 55088.5714 + "H100": { + "throughput_qps": 0.4175, + "latency_mean": 2.3953, + "peak_memory_mb_mean": 55088.5714 + } } }, { @@ -107,9 +113,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.0423, - "latency_mean": 23.5768, - "peak_memory_mb_mean": 66012.2857 + "H100": { + "throughput_qps": 0.0423, + "latency_mean": 23.5768, + "peak_memory_mb_mean": 66012.2857 + } } } ] @@ -152,9 +160,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1214, - "latency_mean": 8.2211, - "peak_memory_mb_mean": 54501.1429 + "H100": { + "throughput_qps": 0.1214, + "latency_mean": 8.2211, + "peak_memory_mb_mean": 54501.1429 + } } } ] @@ -207,9 +217,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.559, - "latency_mean": 1.794, - "peak_memory_mb_mean": 55290.5714 + "H100": { + "throughput_qps": 0.559, + "latency_mean": 1.794, + "peak_memory_mb_mean": 55290.5714 + } } }, { @@ -223,9 +235,11 @@ "max-concurrency": 1, "enable-negative-prompt": true, "baseline": { - "throughput_qps": 0.1918, - "latency_mean": 5.2022, - "peak_memory_mb_mean": 64838.8571 + "H100": { + "throughput_qps": 0.1918, + "latency_mean": 5.2022, + "peak_memory_mb_mean": 64838.8571 + } } } ] @@ -269,9 +283,11 @@ "warmup-num-inference-steps": 20, "enable-negative-prompt": true, "baseline": { - "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159], - "latency_mean": [2.4215, 2.992, 5.9293, 9.656], - "peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821] + "H100": { + "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159], + "latency_mean": [2.4215, 2.992, 5.9293, 9.656], + "peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821] + } } } ] diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index d6d1e0b30..81035a272 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/tests/dfx/perf/tests/test_runner_metadata.py @@ -17,7 +17,7 @@ def test_task_excluded_from_cli_args(): "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "percentile-metrics": "audio_rtf,audio_ttfp", - "baseline": {"mean_audio_rtf": [0.5]}, + "baseline": {"H100": {"mean_audio_rtf": [0.5]}}, } exclude_keys = { "request_rate", @@ -123,35 +123,6 @@ def test_resolve_pytest_marks_rejects_legacy_object_format(): ) -def test_extract_mark_resource_label(): - from tests.dfx.conftest import extract_mark_resource_label - - assert extract_mark_resource_label(None) == "na" - assert ( - extract_mark_resource_label( - [ - {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, - "full_model", - ] - ) - == "H100" - ) - assert ( - extract_mark_resource_label( - [ - { - "hardware_marks": { - "res": {"cuda": "H100", "rocm": "MI325"}, - "num_cards": 2, - } - }, - "diffusion", - ] - ) - == "H100-MI325" - ) - - def test_resource_label_for_filename(): from tests.dfx.conftest import resource_label_for_filename @@ -169,16 +140,14 @@ def test_hardware_json_value(): assert hardware_json_value(None) == "" -def test_extract_configs_resource_label(monkeypatch): - from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label +def test_get_runtime_resource_label(monkeypatch): + from tests.dfx.conftest import get_runtime_resource_label monkeypatch.setattr( "tests.dfx.conftest._read_runtime_device_name", lambda *, device_id=0: "NVIDIA H100 80GB HBM3", ) - get_runtime_resource_label(refresh=True) - assert extract_configs_resource_label([]) == "H100" - get_runtime_resource_label(refresh=True) + assert get_runtime_resource_label(refresh=True) == "H100" monkeypatch.setattr( "tests.dfx.conftest._read_runtime_device_name", lambda *, device_id=0: "Ascend910B2", @@ -336,3 +305,183 @@ def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_pat ("start", tts_server), ("stop", tts_server), ] + + +def test_is_hardware_nested_baseline(): + from tests.dfx.conftest import ( + _RUNTIME_DEVICE_ALIASES, + is_hardware_nested_baseline, + ) + from tests.helpers.mark import get_hardware_mark_list + + hardware_marks = get_hardware_mark_list() + assert hardware_marks + assert {"H100", "L4", "A3", "MI325", "B200"} <= hardware_marks + assert "cuda" not in hardware_marks + assert "full_model" not in hardware_marks + + # Runtime aliases are a full static list (independent of baseline allowlist). + assert "H100" in _RUNTIME_DEVICE_ALIASES + assert "A100" in _RUNTIME_DEVICE_ALIASES + + assert is_hardware_nested_baseline( + { + "H100": {"mean_ttft_ms": [1.0, 2.0], "mean_e2el_ms": [10.0, 20.0]}, + "A3": {"mean_ttft_ms": [0.9, 1.8], "mean_e2el_ms": [9.0, 18.0]}, + } + ) + # Custom metric names under a known hardware label are allowed. + assert is_hardware_nested_baseline({"H100": {"custom_stage_ms": 12.0, "foo_bar": [1.0, 2.0]}}) + assert is_hardware_nested_baseline({"A3": {"request_throughput": 0.5}}) + # Flat metric maps / concurrency-keyed maps are not hardware-nested. + assert not is_hardware_nested_baseline({"mean_ttft_ms": [1.0, 2.0], "mean_e2el_ms": [10.0, 20.0]}) + assert not is_hardware_nested_baseline({"mean_ttft_ms": {"1": 1.0, "32": 2.0}}) + assert not is_hardware_nested_baseline({}) + # Alias-only labels (not [hardware-resource] markers) are rejected for baselines. + assert not is_hardware_nested_baseline({"A100": {"throughput_qps": 1.0}}) + # Unknown top-level hardware label is rejected. + assert not is_hardware_nested_baseline({"UnknownGPU": {"throughput_qps": 1.0}}) + # Empty per-hardware metric map is rejected. + assert not is_hardware_nested_baseline({"H100": {}}) + + +def test_resolve_baseline_for_sweep_keeps_all_hardware_for_one_concurrency(): + from tests.dfx.conftest import resolve_baseline_for_sweep + + baseline = { + "H100": { + "mean_ttft_ms": [96.4, 140.8, 271.9, 362.3, 507.8], + "mean_e2el_ms": [18507.0, 28365.0, 31907.0, 48161.0, 72630.0], + }, + "A3": { + "mean_ttft_ms": [90.0, 130.0, 250.0, 340.0, 480.0], + "mean_e2el_ms": [17000.0, 26000.0, 30000.0, 45000.0, 70000.0], + }, + } + # max_concurrency=[1,4,8,16,32] -> index 4 is concurrency 32 + got = resolve_baseline_for_sweep(baseline, sweep_index=4) + assert got == { + "H100": {"mean_ttft_ms": 507.8, "mean_e2el_ms": 72630.0}, + "A3": {"mean_ttft_ms": 480.0, "mean_e2el_ms": 70000.0}, + } + # First sweep step keeps both hardware buckets too. + got0 = resolve_baseline_for_sweep(baseline, sweep_index=0) + assert set(got0) == {"H100", "A3"} + assert got0["H100"]["mean_ttft_ms"] == 96.4 + assert got0["A3"]["mean_ttft_ms"] == 90.0 + + +def test_resolve_baseline_for_sweep_rejects_flat_baseline(): + from tests.dfx.conftest import resolve_baseline_for_sweep + + with pytest.raises(ValueError, match=r"hardware-nested.*pyproject\.toml"): + resolve_baseline_for_sweep( + {"throughput_qps": [0.4, 0.6], "latency_mean": [1.0, 2.0]}, + sweep_index=1, + ) + + +def test_resolve_baseline_for_sweep_rejects_unknown_hardware_label(): + from tests.dfx.conftest import resolve_baseline_for_sweep + + with pytest.raises(ValueError, match=r"Unknown hardware label\(s\): \['A100'\].*pyproject\.toml"): + resolve_baseline_for_sweep( + {"A100": {"throughput_qps": 1.0}}, + sweep_index=0, + ) + + +def test_resolve_baseline_for_sweep_supports_list_and_scalar_under_hardware(): + from tests.dfx.conftest import resolve_baseline_for_sweep + + # Sweep-aligned lists under each hardware bucket (canonical form). + listed = { + "H100": {"throughput_qps": [0.4, 0.6, 0.8], "latency_mean": [1.0, 2.0, 3.0]}, + "A3": {"throughput_qps": [0.5, 0.7, 0.9], "latency_mean": [0.9, 1.8, 2.7]}, + } + assert resolve_baseline_for_sweep(listed, sweep_index=1) == { + "H100": {"throughput_qps": 0.6, "latency_mean": 2.0}, + "A3": {"throughput_qps": 0.7, "latency_mean": 1.8}, + } + + # Scalars under hardware stay as-is (single-concurrency cases). + scalar = {"H100": {"throughput_qps": 0.5}} + assert resolve_baseline_for_sweep(scalar, sweep_index=0) == {"H100": {"throughput_qps": 0.5}} + # Custom metric names are preserved. + custom = {"H100": {"custom_stage_ms": [10.0, 20.0]}} + assert resolve_baseline_for_sweep(custom, sweep_index=1) == {"H100": {"custom_stage_ms": 20.0}} + assert resolve_baseline_for_sweep(None) == {} + assert resolve_baseline_for_sweep({}) == {} + + +def test_resolve_baseline_value_errors(): + from tests.dfx.conftest import resolve_baseline_value + + with pytest.raises(ValueError, match="sweep_index"): + resolve_baseline_value([1.0, 2.0], sweep_index=None) + with pytest.raises(IndexError): + resolve_baseline_value([1.0], sweep_index=1) + with pytest.raises(TypeError, match="not supported"): + resolve_baseline_value({"1": 0.4}, sweep_index=0) + + +def test_diffusion_build_run_params_resolves_baseline_per_sweep(tmp_path, monkeypatch): + """``_build_run_params`` / ``_iter_sweep_runs`` narrow list baselines per concurrency.""" + import importlib + import sys + + cfg = tmp_path / "mini_diffusion_perf.json" + cfg.write_text( + json.dumps( + [ + { + "test_name": "test_mini", + "server_type": "vllm-omni", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, + "diffusion", + "full_model", + ], + "server_params": {"model": "m/mini"}, + "benchmark_params": [{"name": "p0", "num-prompts": 1, "max-concurrency": 1}], + } + ] + ), + encoding="utf-8", + ) + monkeypatch.setattr(sys, "argv", ["pytest", "--test-config-file", str(cfg)]) + sys.modules.pop("tests.dfx.perf.scripts.run_diffusion_benchmark", None) + from tests.dfx.perf.scripts import run_diffusion_benchmark as rdb + + importlib.reload(rdb) + + params = { + "name": "c_sweep", + "dataset": "random", + "task": "t2i", + "num-prompts": [8, 16, 32], + "max-concurrency": [1, 8, 32], + "baseline": { + "H100": {"throughput_qps": [0.1, 0.2, 0.3], "latency_mean": [10.0, 20.0, 30.0]}, + "A3": {"throughput_qps": [0.11, 0.22, 0.33], "latency_mean": [9.0, 19.0, 29.0]}, + }, + } + run32 = rdb._build_run_params( + params, + num_prompts=32, + max_concurrency=32, + request_rate="inf", + sweep_index=2, + ) + assert run32["max-concurrency"] == 32 + assert run32["baseline"] == { + "H100": {"throughput_qps": 0.3, "latency_mean": 30.0}, + "A3": {"throughput_qps": 0.33, "latency_mean": 29.0}, + } + + sweeps = rdb._iter_sweep_runs(params) + assert len(sweeps) == 3 + assert sweeps[2]["params"]["max-concurrency"] == 32 + assert sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"] == 0.3 + assert "A3" in sweeps[2]["params"]["baseline"] + assert not isinstance(sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"], list) diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index 91037aab6..f398e3271 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -39,12 +39,14 @@ "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 141.2791 - ], - "median_audio_rtf": [ - 0.1591 - ] + "H100": { + "median_audio_ttfp_ms": [ + 141.2791 + ], + "median_audio_rtf": [ + 0.1591 + ] + } } }, { @@ -67,21 +69,23 @@ "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 226.747, - 647.1272, - 4777.2714 - ], - "median_audio_rtf": [ - 0.3313, - 0.3642, - 1.2282 - ], - "audio_throughput": [ - 24.2296, - 40.3552, - 43.2572 - ] + "H100": { + "median_audio_ttfp_ms": [ + 226.747, + 647.1272, + 4777.2714 + ], + "median_audio_rtf": [ + 0.3313, + 0.3642, + 1.2282 + ], + "audio_throughput": [ + 24.2296, + 40.3552, + 43.2572 + ] + } } }, { @@ -102,9 +106,11 @@ "seed_tts_wer_eval": true, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_audio_rtf": [ - 0.45 - ] + "H100": { + "mean_audio_rtf": [ + 0.45 + ] + } } } ] @@ -153,12 +159,14 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 46.8434 - ], - "median_audio_rtf": [ - 0.1367 - ] + "H100": { + "median_audio_ttfp_ms": [ + 46.8434 + ], + "median_audio_rtf": [ + 0.1367 + ] + } } }, { @@ -186,21 +194,23 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 75.3341, - 713.7727, - 5941.7179 - ], - "median_audio_rtf": [ - 0.1859, - 0.3148, - 1.1662 - ], - "audio_throughput": [ - 38.3858, - 55.3523, - 68.7032 - ] + "H100": { + "median_audio_ttfp_ms": [ + 75.3341, + 713.7727, + 5941.7179 + ], + "median_audio_rtf": [ + 0.1859, + 0.3148, + 1.1662 + ], + "audio_throughput": [ + 38.3858, + 55.3523, + 68.7032 + ] + } } }, { @@ -226,9 +236,11 @@ "seed_tts_wer_eval": true, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_audio_rtf": [ - 0.35 - ] + "H100": { + "mean_audio_rtf": [ + 0.35 + ] + } } }, { @@ -252,15 +264,17 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 56.746 - ], - "median_audio_rtf": [ - 0.1571 - ], - "audio_throughput": [ - 11.1226 - ] + "H100": { + "median_audio_ttfp_ms": [ + 56.746 + ], + "median_audio_rtf": [ + 0.1571 + ], + "audio_throughput": [ + 11.1226 + ] + } } }, { @@ -283,12 +297,14 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 46.3479 - ], - "median_audio_rtf": [ - 0.1352 - ] + "H100": { + "median_audio_ttfp_ms": [ + 46.3479 + ], + "median_audio_rtf": [ + 0.1352 + ] + } } }, { @@ -315,21 +331,23 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 75.5967, - 587.5877, - 4960.3726 - ], - "median_audio_rtf": [ - 0.1871, - 0.3143, - 1.1703 - ], - "audio_throughput": [ - 38.9378, - 55.1846, - 71.5708 - ] + "H100": { + "median_audio_ttfp_ms": [ + 75.5967, + 587.5877, + 4960.3726 + ], + "median_audio_rtf": [ + 0.1871, + 0.3143, + 1.1703 + ], + "audio_throughput": [ + 38.9378, + 55.1846, + 71.5708 + ] + } } }, { @@ -352,15 +370,17 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 55.5029 - ], - "median_audio_rtf": [ - 0.1508 - ], - "audio_throughput": [ - 9.4885 - ] + "H100": { + "median_audio_ttfp_ms": [ + 55.5029 + ], + "median_audio_rtf": [ + 0.1508 + ], + "audio_throughput": [ + 9.4885 + ] + } } } ] @@ -402,12 +422,14 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 82.1886 - ], - "median_audio_rtf": [ - 0.1241 - ] + "H100": { + "median_audio_ttfp_ms": [ + 82.1886 + ], + "median_audio_rtf": [ + 0.1241 + ] + } } }, { @@ -429,15 +451,17 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 160.9168 - ], - "median_audio_rtf": [ - 0.1907 - ], - "audio_throughput": [ - 39.2332 - ] + "H100": { + "median_audio_ttfp_ms": [ + 160.9168 + ], + "median_audio_rtf": [ + 0.1907 + ], + "audio_throughput": [ + 39.2332 + ] + } } }, { @@ -459,15 +483,17 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 100.1265 - ], - "median_audio_rtf": [ - 0.1437 - ], - "audio_throughput": [ - 11.7605 - ] + "H100": { + "median_audio_ttfp_ms": [ + 100.1265 + ], + "median_audio_rtf": [ + 0.1437 + ], + "audio_throughput": [ + 11.7605 + ] + } } } ] diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json index 12afdc132..e5163e5bb 100644 --- a/tests/dfx/perf/tests/test_voxcpm2.json +++ b/tests/dfx/perf/tests/test_voxcpm2.json @@ -37,12 +37,14 @@ "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 137.645 - ], - "median_audio_rtf": [ - 0.0497 - ] + "H100": { + "median_audio_ttfp_ms": [ + 137.645 + ], + "median_audio_rtf": [ + 0.0497 + ] + } } }, { @@ -61,15 +63,17 @@ "seed_tts_locale": "en", "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 475.3357 - ], - "median_audio_rtf": [ - 0.1243 - ], - "audio_throughput": [ - 36.6358 - ] + "H100": { + "median_audio_ttfp_ms": [ + 475.3357 + ], + "median_audio_rtf": [ + 0.1243 + ], + "audio_throughput": [ + 36.6358 + ] + } } }, { @@ -90,9 +94,11 @@ "seed_tts_wer_eval": true, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_audio_rtf": [ - 0.45 - ] + "H100": { + "mean_audio_rtf": [ + 0.45 + ] + } } }, { @@ -114,12 +120,14 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 79.6927 - ], - "median_audio_rtf": [ - 0.0423 - ] + "H100": { + "median_audio_ttfp_ms": [ + 79.6927 + ], + "median_audio_rtf": [ + 0.0423 + ] + } } }, { @@ -141,15 +149,17 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 262.6212 - ], - "median_audio_rtf": [ - 0.0753 - ], - "audio_throughput": [ - 103.5843 - ] + "H100": { + "median_audio_ttfp_ms": [ + 262.6212 + ], + "median_audio_rtf": [ + 0.0753 + ], + "audio_throughput": [ + 103.5843 + ] + } } }, { @@ -171,15 +181,17 @@ }, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "median_audio_ttfp_ms": [ - 80.1095 - ], - "median_audio_rtf": [ - 0.0479 - ], - "audio_throughput": [ - 17.2918 - ] + "H100": { + "median_audio_ttfp_ms": [ + 80.1095 + ], + "median_audio_rtf": [ + 0.0479 + ], + "audio_throughput": [ + 17.2918 + ] + } } }, { @@ -203,9 +215,11 @@ "seed_tts_wer_eval": true, "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_audio_rtf": [ - 0.35 - ] + "H100": { + "mean_audio_rtf": [ + 0.35 + ] + } } } ] diff --git a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json index 52441dae4..5cea0c013 100644 --- a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json +++ b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json @@ -42,9 +42,11 @@ } ], "baseline": { - "throughput_qps": 0.0434, - "latency_mean": 23.0654, - "peak_memory_mb_mean": 76360.0 + "H100": { + "throughput_qps": 0.0434, + "latency_mean": 23.0654, + "peak_memory_mb_mean": 76360.0 + } } } ] @@ -97,9 +99,11 @@ } ], "baseline": { - "throughput_qps": 0.0606, - "latency_mean": 16.4928, - "peak_memory_mb_mean": 47164.8 + "H100": { + "throughput_qps": 0.0606, + "latency_mean": 16.4928, + "peak_memory_mb_mean": 47164.8 + } } }, { @@ -122,9 +126,11 @@ } ], "baseline": { - "throughput_qps": 0.0124, - "latency_mean": 80.7784, - "peak_memory_mb_mean": 55374.3429 + "H100": { + "throughput_qps": 0.0124, + "latency_mean": 80.7784, + "peak_memory_mb_mean": 55374.3429 + } } } ] diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 72b409de5..4a2c65cbb 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -90,6 +90,35 @@ class TestAttentionSpec: with pytest.raises(ValueError, match=match): AttentionSpec(**spec) + def test_block_sparse_defaults_applied_when_backend_selected(self): + spec = AttentionSpec(backend="RAINFUSION_ATTN") + assert spec.block_sparse.sparsity == 0.8 + assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0} + + def test_block_sparse_skip_layers_selector_expanded(self): + spec = AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse={"sparsity": 0.9, "start_step": 12, "skip_layers": "0-2,38"}, + ) + assert spec.block_sparse.skip_layer_indices == {0, 1, 2, 38} + assert spec.backend_kwargs() == { + "sparsity": 0.9, + "start_step": 12, + "skip_layers": [0, 1, 2, 38], + } + + def test_block_sparse_rejected_on_dense_backend(self): + with pytest.raises(ValueError, match="block_sparse is only supported by"): + AttentionSpec(backend="FLASH_ATTN", block_sparse={"sparsity": 0.8}) + + @pytest.mark.parametrize( + "block_sparse", + [{"sparsity": 1.5}, {"start_step": -1}], + ) + def test_block_sparse_invalid_values(self, block_sparse): + with pytest.raises(ValueError): + AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse) + class TestAttentionConfig: def test_empty_config(self): diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py new file mode 100644 index 000000000..de1adf8c7 --- /dev/null +++ b/tests/diffusion/attention/test_cudnn_attn.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import nullcontext + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): + observed = {} + + def fake_sdpa(query, key, value, **kwargs): + observed.update(query=query, key=key, value=value, kwargs=kwargs) + return query + + monkeypatch.setattr( + "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel", + lambda _backends: nullcontext(), + ) + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=4, + softmax_scale=0.5, + ) + query = torch.randn(1, 8, 2, 4) + key = torch.randn_like(query) + value = torch.randn_like(query) + + output = impl.forward_cuda( + query, + key, + value, + AttentionMetadata(extra={"valid_kv_length": 5}), + ) + + assert output.shape == query.shape + assert observed["query"].shape == (1, 2, 8, 4) + assert observed["key"].shape == (1, 2, 5, 4) + assert observed["value"].shape == (1, 2, 5, 4) + assert observed["kwargs"]["attn_mask"] is None + + +def test_cudnn_rejects_invalid_valid_kv_length(): + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=4, + softmax_scale=0.5, + ) + query = torch.randn(1, 8, 2, 4) + + with pytest.raises(ValueError, match="valid_kv_length"): + impl.forward_cuda( + query, + query, + query, + AttentionMetadata(extra={"valid_kv_length": 9}), + ) diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py new file mode 100644 index 000000000..e88a973a9 --- /dev/null +++ b/tests/diffusion/attention/test_rainfusion_plan.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2. + +The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip: +1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row +kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the +first can run sparse: rf_v2 tiles the sequence as a whole but pools its block +mask from the video and prefix segments separately, so the two only agree on the +block count when the video segment ends on a boundary. +""" + +import dataclasses + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout +from vllm_omni.diffusion.attention.backends.rainfusion_attn import ( + _BLOCK_SIZE, + RainFusionAttentionImpl, + RainFusionPlan, +) +from vllm_omni.platforms import current_omni_platform + +PREFIX_ROWS = 710 # 14 text rows + 696 audio rows +ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks +MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows + + +def make_impl(**backend_kwargs): + return RainFusionAttentionImpl( + num_heads=8, + head_size=128, + softmax_scale=128**-0.5, + prefix="transformer_blocks.0.attn", + qkv_layout="BSND", + backend_kwargs={"sparsity": 0.8, **backend_kwargs}, + ) + + +def make_metadata(grid, prefix_len=PREFIX_ROWS, max_seqlen_q=None): + video_rows = grid[0] * grid[1] * grid[2] + return AttentionMetadata( + extra={"max_seqlen_q": prefix_len + video_rows if max_seqlen_q is None else max_seqlen_q}, + video_layout=VideoTokenLayout(prefix_len=prefix_len, latent_grid=grid), + ) + + +def test_block_aligned_video_segment_runs_sparse(): + plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID)) + + assert plan is not None + assert plan.prefix_len == PREFIX_ROWS + assert plan.used_len == PREFIX_ROWS + 59520 + assert plan.latent_shape == list(ALIGNED_GRID) + + +def test_misaligned_video_segment_stays_dense(): + # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask: + # pad keys would take a share of every real query's softmax denominator. + assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None + + +@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900]) +def test_prefix_length_does_not_affect_alignment(prefix_len): + # Only the video segment has to land on a block boundary; rf_v2 pools the + # prefix separately and keeps every one of its blocks. + plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID, prefix_len=prefix_len)) + + assert plan is not None + assert plan.prefix_len == prefix_len + + +def test_sparsity_zero_never_resolves_a_plan(): + assert make_impl(sparsity=0.0)._resolve_plan(make_metadata(ALIGNED_GRID)) is None + + +def test_missing_video_layout_falls_back_to_dense(): + assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 60230})) is None + + +def test_missing_max_seqlen_falls_back_to_dense(): + metadata = AttentionMetadata( + video_layout=VideoTokenLayout(prefix_len=PREFIX_ROWS, latent_grid=ALIGNED_GRID), + ) + + assert make_impl()._resolve_plan(metadata) is None + + +def test_video_segment_must_be_the_tail_of_packed_document_zero(): + metadata = make_metadata(ALIGNED_GRID, max_seqlen_q=PREFIX_ROWS + 59520 + 128) + + assert make_impl()._resolve_plan(metadata) is None + + +@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)]) +def test_short_video_stays_dense(grid): + assert make_impl()._resolve_plan(make_metadata(grid)) is None + + +@pytest.mark.parametrize("qkv_layout", ["BNSD", "BSH"]) +def test_explicitly_wrong_layout_is_rejected(qkv_layout): + with pytest.raises(ValueError, match="BSND"): + RainFusionAttentionImpl( + num_heads=8, + head_size=128, + softmax_scale=128**-0.5, + qkv_layout=qkv_layout, + backend_kwargs={"sparsity": 0.8}, + ) + + +def test_undeclared_layout_stays_dense(): + # An undeclared layout is not an error -- the layer keeps working, just + # densely, and the fallback sees the same absent layout plain FLASH_ATTN would. + impl = RainFusionAttentionImpl( + num_heads=8, + head_size=128, + softmax_scale=128**-0.5, + prefix="transformer_blocks.0.attn", + backend_kwargs={"sparsity": 0.8}, + ) + + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None + assert impl.dense_fallback.qkv_layout is None + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.") +def test_fully_populated_mask_reproduces_dense_attention(): + """At sparsity=0 every key block is kept, so rf_v2 must match dense attention. + + This is the guard on the geometry the backend hands the kernel: a wrong + prefix length or latent grid still produces plausible output under sparse + selection, but shows up here as a mismatch against SDPA. + """ + torch.manual_seed(0) + prefix_len, grid = 384, (4, 16, 64) # 4096 video rows = 32 blocks + video_len = grid[0] * grid[1] * grid[2] + used = prefix_len + video_len + heads, head_dim = 4, 128 + q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3)) + + impl = make_impl() + impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0) + plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid)) + out = impl._forward_sparse_npu(q, k, v, plan) + + reference = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + scale=impl.softmax_scale, + ).transpose(1, 2) + error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean() + assert error < 2e-3, f"mean relative error {error:.4%} against dense attention" diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py index a2f855596..76daf9d3b 100644 --- a/tests/diffusion/attention/test_trtllm_attn.py +++ b/tests/diffusion/attention/test_trtllm_attn.py @@ -8,6 +8,7 @@ import torch import torch.nn.functional as F from vllm_omni.diffusion.attention.backends import trtllm_attn as tg +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] @@ -142,6 +143,9 @@ def test_denoise_progress_mixin(monkeypatch): _Pipe().record_denoise_step(4) assert ctx.denoise_step_idx == 4 and ctx.denoise_timestep == pytest.approx(0.25) + _Pipe().record_denoise_step(5, normalized_timestep=0.7) + assert ctx.denoise_step_idx == 5 and ctx.denoise_timestep == pytest.approx(0.7) + @pytest.mark.parametrize( "bk, field", @@ -187,6 +191,7 @@ def test_mask_free_allowlist_gates_trtllm_default(): assert meta("WanPipeline").attention_mask_free is True assert meta("WanVACEPipeline").attention_mask_free is True + assert meta("MiniMaxH3Pipeline").attention_mask_free is True assert meta("FluxPipeline").attention_mask_free is False assert meta(None).attention_mask_free is False @@ -204,6 +209,91 @@ def test_masked_layer_raises(): _impl().forward_cuda(q, k, v, _Meta()) +def test_packed_metadata_trims_padding_mask(monkeypatch): + b, s, h, d = 1, 8, 8, 128 + q, k, v = (torch.zeros(b, s, h, d) for _ in range(3)) + mask = torch.tensor([[True, True, True, True, True, True, False, False]]) + cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32) + metadata = AttentionMetadata( + attn_mask=mask, + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": 6, + "max_seqlen_k": 6, + }, + ) + captured = {} + + def fake_attention(**kwargs): + captured.update(kwargs) + return torch.zeros_like(kwargs["query"]) + + monkeypatch.setattr(tg, "HAS_FLASHINFER", True) + monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention) + monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0))) + + out = _impl().forward_cuda(q, k, v, metadata) + + assert out.shape == q.shape + assert captured["query"].shape[0] == 6 + assert captured["key"].shape[0] == 6 + assert captured["value"].shape[0] == 6 + assert captured["batch_size"] == 1 + assert captured["seq_lens"].tolist() == [6] + assert captured["cum_seq_lens_q"].tolist() == [0, 6] + assert captured["cum_seq_lens_kv"].tolist() == [0, 6] + assert captured["max_q_len"] == 6 + assert captured["max_kv_len"] == 6 + assert torch.count_nonzero(out[:, 6:]) == 0 + + +def test_packed_metadata_rejects_non_prefix_mask(monkeypatch): + b, s, h, d = 1, 8, 8, 128 + q, k, v = (torch.zeros(b, s, h, d) for _ in range(3)) + cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32) + metadata = AttentionMetadata( + attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]), + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": 6, + "max_seqlen_k": 6, + }, + ) + monkeypatch.setattr(tg, "HAS_FLASHINFER", True) + + with pytest.raises(ValueError, match="prefix-valid"): + _impl().forward_cuda(q, k, v, metadata) + + +def test_packed_metadata_must_be_complete(): + b, s, h, d = 1, 8, 8, 128 + q, k, v = (torch.zeros(b, s, h, d) for _ in range(3)) + metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, s], dtype=torch.int32)}) + + with pytest.raises(ValueError, match="Incomplete packed TRTLLM attention metadata"): + _impl().forward_cuda(q, k, v, metadata) + + +def test_packed_metadata_must_cover_inputs(monkeypatch): + b, s, h, d = 1, 8, 8, 128 + q, k, v = (torch.zeros(b, s, h, d) for _ in range(3)) + cu_seqlens = torch.tensor([0, 6], dtype=torch.int32) + metadata = AttentionMetadata( + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": 6, + "max_seqlen_k": 6, + } + ) + monkeypatch.setattr(tg, "HAS_FLASHINFER", True) + + with pytest.raises(ValueError, match="must cover all Q/K/V tokens"): + _impl().forward_cuda(q, k, v, metadata) + + def test_propagate_calibration_045_schema(monkeypatch): import types @@ -295,6 +385,30 @@ def test_bf16_dense_matches_sdpa(): assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high" +@requires_trtllm_attn +def test_bf16_packed_padding_matches_sdpa(): + torch.manual_seed(0) + b, s, used, h, d = 1, 256, 192, 8, 128 + scale = 1.0 / math.sqrt(d) + q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda") + metadata = AttentionMetadata( + attn_mask=torch.arange(s, device="cuda")[None] < used, + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": used, + "max_seqlen_k": used, + }, + ) + + out = _impl().forward_cuda(q, k, v, metadata).float() + ref_used = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale) + used_rel = (out[:, :used] - ref_used).abs().mean() / ref_used.abs().mean() + assert used_rel < 0.01, f"BF16 packed valid-token rel err {used_rel:.4f} too high" + assert torch.count_nonzero(out[:, used:]) == 0 + + requires_sage = pytest.mark.skipif( not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None), reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)", @@ -321,6 +435,47 @@ def test_sage_quant_matches_sdpa(dtype_qk): assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high" +@requires_sage +def test_sage_packed_non_aligned_length_matches_sdpa(): + torch.manual_seed(0) + b, s, used, h, d = 1, 208, 194, 8, 128 + scale = 1.0 / math.sqrt(d) + q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda") + metadata = AttentionMetadata( + attn_mask=torch.arange(s, device="cuda")[None] < used, + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": used, + "max_seqlen_k": used, + }, + ) + + out = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}).forward_cuda(q, k, v, metadata) + assert torch.isfinite(out).all() + ref = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale) + rel = (out[:, :used].float() - ref).abs().mean() / ref.abs().mean() + assert rel < 0.2, f"SAGE packed valid-token rel err {rel:.4f} too high" + assert torch.count_nonzero(out[:, used:]) == 0 + + +@requires_sage +def test_sage_short_sequence_uses_dense_kernel(): + torch.manual_seed(0) + b, s, h, d = 1, 14, 8, 128 + scale = 1.0 / math.sqrt(d) + q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + + impl = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}) + impl._sage_quantize_fn = lambda *args, **kwargs: pytest.fail("short sequence must not use SAGE") + out = impl.forward_cuda(q, k, v) + ref = _sdpa_ref(q, k, v, scale) + rel = (out.float() - ref).abs().mean() / ref.abs().mean() + assert torch.isfinite(out).all() + assert rel < 0.01, f"SAGE short-sequence dense fallback rel err {rel:.4f} too high" + + @requires_trtllm_attn def test_skip_tiny_threshold_matches_dense(): torch.manual_seed(0) diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index bb7e9f2c9..1137cfbc0 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -86,9 +86,35 @@ class TestCacheDiTBackend: # Verify cache-dit was enabled assert backend.enabled is True - assert backend._refresh_func is not None + assert backend._refresh_funcs mock_cache_dit.enable_cache.assert_called_once() + @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") + @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") + def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter): + """Combined pipelines enable and refresh Cache-DiT for every DiT.""" + pipeline = Mock() + pipeline.__class__.__name__ = "MiniMaxH3Pipeline" + pipeline._dit_modules = ["transformer", "transformers_ref"] + pipeline.transformer = Mock() + pipeline.transformers_ref = Mock() + for transformer in (pipeline.transformer, pipeline.transformers_ref): + transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={"blocks": ForwardPattern.Pattern_3} + ) + + backend = CacheDiTBackend({"Fn_compute_blocks": 2}) + backend.enable(pipeline) + backend.refresh(pipeline, num_inference_steps=20) + + assert len(backend._refresh_funcs) == 2 + assert mock_cache_dit.enable_cache.call_count == 2 + assert mock_cache_dit.refresh_context.call_count == 2 + assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == { + pipeline.transformer, + pipeline.transformers_ref, + } + @patch("vllm_omni.diffusion.cache.cachedit.backend.logger") @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") def test_enable_without_model_adapter_uses_cache_dit_registry(self, mock_cache_dit, mock_logger): @@ -106,7 +132,7 @@ class TestCacheDiTBackend: backend.enable(pipeline) assert backend.enabled is True - assert backend._refresh_func is not None + assert backend._refresh_funcs assert mock_cache_dit.enable_cache.call_args.args[0] is pipeline.transformer mock_logger.info.assert_any_call( "Transformer %s does not declare _cache_dit_adapter_config; " @@ -138,7 +164,7 @@ class TestCacheDiTBackend: assert isinstance(exc_info.value.__cause__, ValueError) assert backend.enabled is False - assert backend._refresh_func is None + assert backend._refresh_funcs == [] @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") @@ -245,7 +271,7 @@ class TestCacheDiTBackend: backend.enable(mock_pipeline) assert backend.enabled is True - assert backend._refresh_func is not None + assert backend._refresh_funcs mock_block_adapter.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["transformer"] is mock_pipeline.model @@ -285,7 +311,7 @@ class TestCacheDiTBackend: backend.enable(mock_pipeline) assert backend.enabled is True - assert backend._refresh_func is not None + assert backend._refresh_funcs mock_block_adapter.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["transformer"] is mock_pipeline.transformer diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 694f8f8e9..71125684b 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/tests/diffusion/layers/test_norm.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers.""" +import sys +import types + import pytest import torch from pytest_mock import MockerFixture @@ -299,6 +302,44 @@ def test_rmsnorm_numerical_correctness(): torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) +def test_rmsnorm_dtype_matches_parameter_and_fp32_accumulation_contract(): + """dtype controls gamma storage; reduced-precision inputs still use fp32 accumulation.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + eps = 1e-6 + norm = RMSNorm(hidden_size, eps=eps, dtype=torch.bfloat16) + x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16) + + assert norm.weight.dtype == torch.bfloat16 + + x_fp32 = x.float() + expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype) + + torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0) + + +def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.MonkeyPatch): + """The NPU fused op receives the bf16 checkpoint gamma used by H3.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + captured: dict[str, torch.Tensor] = {} + + def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float): + captured["gamma"] = gamma + captured["epsilon"] = torch.tensor(epsilon) + return (x,) + + monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm)) + norm = RMSNorm(64, eps=1e-5, dtype=torch.bfloat16) + x = torch.randn(2, 4, 64, dtype=torch.bfloat16) + + assert norm.forward_npu(x) is x + assert captured["gamma"] is norm.weight + assert captured["gamma"].dtype == torch.bfloat16 + assert captured["epsilon"].item() == pytest.approx(1e-5) + + # ── RMSNorm compile-path regression tests ── diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py index 29fcf1804..e1238ba44 100644 --- a/tests/diffusion/layers/test_rope_broadcast.py +++ b/tests/diffusion/layers/test_rope_broadcast.py @@ -10,10 +10,13 @@ forward_hip, and apply_rotary_emb_mindiesd paths. from __future__ import annotations +import sys +import types + import pytest import torch -from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -234,3 +237,33 @@ class TestRotaryEmbeddingNativeShapeRegression: output = rope.forward_native(x, cos, sin) assert output.shape == x.shape assert not torch.isnan(output).any() + + +def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None: + """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout.""" + captured: dict[str, object] = {} + + def rotary_position_embedding(x, cos, sin, **kwargs): + captured.update(x=x, cos=cos, sin=sin, kwargs=kwargs) + return x + + monkeypatch.setitem( + sys.modules, + "mindiesd", + types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding), + ) + x = torch.randn(11, 3, 96, dtype=torch.bfloat16) + cos = torch.randn(11, 96, dtype=torch.bfloat16) + sin = torch.randn(11, 96, dtype=torch.bfloat16) + + actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False) + + assert actual is x + assert captured["x"] is x + assert captured["cos"] is cos + assert captured["sin"] is sin + assert captured["kwargs"] == { + "rotated_mode": "rotated_half", + "head_first": False, + "fused": True, + } diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 95c379d7b..70f00f4b7 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from unittest.mock import Mock @@ -24,12 +25,343 @@ def test_pipeline_import_registry_and_component_discovery(): "pipeline_minimax_h3", "MiniMaxH3Pipeline", ) + assert _DIFFUSION_MODELS["MiniMaxH3ModularPipeline"] == _DIFFUSION_MODELS["MiniMaxH3Pipeline"] assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func" - assert MiniMaxH3Pipeline._dit_modules == ["transformer"] + assert ( + _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3ModularPipeline"] == _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] + ) + assert MiniMaxH3Pipeline._dit_modules == ["transformer", "transformers_ref"] assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"] assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"] +def _write_partition_index(path, *, partition, tasks): + path.mkdir(parents=True) + (path / "model_index.json").write_text( + json.dumps( + { + "_class_name": "MiniMaxH3Pipeline", + "_diffusers_version": "0.32.2", + "_minimax_h3": { + "partition": partition, + "tasks": tasks, + "sigma_shift_scales": {"video": 12.0, "audio": 3.0}, + }, + } + ), + encoding="utf-8", + ) + + +def test_modular_diffusers_index_is_resolved_generically(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name + from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model + from vllm_omni.entrypoints.utils import resolve_model_config_path + + (tmp_path / "modular_model_index.json").write_text( + json.dumps( + { + "_class_name": "MiniMaxH3ModularPipeline", + "_diffusers_version": "0.36.0.dev0", + } + ), + encoding="utf-8", + ) + + assert is_diffusion_model(str(tmp_path)) + assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline" + assert resolve_model_config_path(str(tmp_path)) is None + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + assert config.model_class_name == "MiniMaxH3ModularPipeline" + assert config.supports_multimodal_inputs + assert config.max_multimodal_image_inputs == 9 + assert config.supports_mixed_reference_inputs + + +@pytest.mark.parametrize( + ("task_type", "partition"), + [ + (None, "combined"), + ("auto", "combined"), + ("t2va", "fl2va"), + ("fl2va", "fl2va"), + ("ref2va", "ref2va"), + ], +) +def test_startup_task_selects_weight_partition(task_type, partition): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task + + assert _minimax_h3_partition_for_task(task_type) == partition + + +def test_startup_auto_task_uses_explicit_local_partition(tmp_path): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task + + for directory, partition, tasks in ( + ("FL2VA", "fl2va", ["t2va", "fl2va"]), + ("Ref2VA", "ref2va", ["ref2va"]), + ): + model_path = tmp_path / directory + _write_partition_index(model_path, partition=partition, tasks=tasks) + assert _minimax_h3_partition_for_task(None, str(model_path)) == partition + assert _minimax_h3_partition_for_task("auto", str(model_path)) == partition + assert _minimax_h3_partition_for_task("combined", str(model_path)) == "combined" + + +def test_startup_task_rejects_unsupported_value(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _minimax_h3_partition_for_task + + with pytest.raises(ValueError, match="task_type must be one of"): + _minimax_h3_partition_for_task("unsupported") + + +def test_combined_task_inference_and_transformer_routing(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "combined" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline.transformer = torch.nn.Identity() + pipeline.transformers_ref = torch.nn.Linear(1, 1) + assert pipeline._resolve_task(None, {}) == "t2va" + assert pipeline._resolve_task(None, {"image": object()}) == "fl2va" + assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va" + assert pipeline._resolve_task(None, {"video": object()}) == "ref2va" + pipeline.partition = "ref2va" + pipeline.supported_tasks = frozenset({"ref2va"}) + assert pipeline._resolve_task(None, {"image": object()}) == "ref2va" + + pipeline.partition = "combined" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + assert pipeline._transformer_for_task("fl2va") is pipeline.transformer + assert pipeline._transformer_for_task("ref2va") is pipeline.transformers_ref + + +@pytest.mark.parametrize( + ( + "task_type", + "expected_partition", + "expected_dits", + "component_partition", + "source_partitions", + "expected_tasks", + ), + [ + (None, "combined", 2, "FL2VA", ["FL2VA", "Ref2VA"], {"t2va", "fl2va", "ref2va"}), + ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}), + ("ref2va", "ref2va", 1, "Ref2VA", ["Ref2VA"], {"ref2va"}), + ], +) +def test_pipeline_loads_task_selected_dits_and_shared_components_once( + monkeypatch, + tmp_path, + task_type, + expected_partition, + expected_dits, + component_partition, + source_partitions, + expected_tasks, +): + from vllm_omni.diffusion.data import ( + DiffusionParallelConfig, + OmniDiffusionConfig, + ) + from vllm_omni.diffusion.models.minimax_h3 import ( + pipeline_minimax_h3 as pipeline_module, + ) + + _write_partition_index( + tmp_path / "FL2VA", + partition="fl2va", + tasks=["t2va", "fl2va"], + ) + _write_partition_index( + tmp_path / "Ref2VA", + partition="ref2va", + tasks=["ref2va"], + ) + for partition_name in ("FL2VA", "Ref2VA"): + for component in ( + "transformer", + "tokenizer", + "processor", + "text_encoder", + "video_vae", + "audio_vae", + ): + (tmp_path / partition_name / component).mkdir() + + created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} + + class FakeModule(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + class FakeDiT(FakeModule): + def __init__(self, *args, **kwargs): + super().__init__() + created["dit"].append((args, kwargs)) + + def component_factory(name): + def create(path, *args, **kwargs): + created[name].append(str(path)) + return FakeModule() + + return create + + tokenizer_calls = [] + processor_calls = [] + tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast) + tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: ( + tokenizer_calls.append((args, kwargs)) or object() + ) + processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor) + processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: ( + processor_calls.append((args, kwargs)) or object() + ) + monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT) + monkeypatch.setattr( + pipeline_module, + "MiniMaxH3Qwen3VLEncoder", + component_factory("text_encoder"), + ) + monkeypatch.setattr( + pipeline_module, + "MiniMaxH3VideoVAE", + component_factory("video_vae"), + ) + monkeypatch.setattr( + pipeline_module, + "MiniMaxH3AudioVAE", + component_factory("audio_vae"), + ) + monkeypatch.setattr( + pipeline_module, + "Qwen2TokenizerFast", + tokenizer_cls, + ) + monkeypatch.setattr( + pipeline_module, + "Qwen3VLProcessor", + processor_cls, + ) + monkeypatch.setattr( + pipeline_module, + "_dit_rank_world", + lambda: (None, 0, 1), + ) + monkeypatch.setattr( + pipeline_module, + "get_local_device", + lambda: torch.device("cpu"), + ) + download_calls = [] + + def fake_download(**kwargs): + download_calls.append(kwargs) + return str(tmp_path) + + monkeypatch.setattr( + pipeline_module, + "download_weights_from_hf_specific", + fake_download, + ) + + od_config = OmniDiffusionConfig( + model="MiniMaxAI/MiniMax-H3", + revision=None, + task_type=task_type, + quantization_config=None, + parallel_config=DiffusionParallelConfig( + cfg_parallel_size=1, + text_encoder_tp_size=1, + ), + ) + pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config) + + assert pipeline.partition == expected_partition + assert pipeline.supported_tasks == expected_tasks + assert len(created["dit"]) == expected_dits + component_path = tmp_path / component_partition + assert created["text_encoder"] == [str(component_path / "text_encoder")] + assert created["video_vae"] == [str(component_path / "video_vae")] + assert created["audio_vae"] == [str(component_path / "audio_vae")] + assert len(tokenizer_calls) == 1 + assert len(processor_calls) == 1 + expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"] + assert pipeline._dit_modules == expected_dit_modules + assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2) + if expected_partition == "ref2va": + assert pipeline._transformer_for_task("ref2va") is pipeline.transformer + assert [source.model_or_path for source in pipeline.weights_sources] == [ + str(tmp_path / partition_name) for partition_name in source_partitions + ] + expected_patterns = ( + pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS + if expected_partition == "combined" + else pipeline_module.MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[expected_partition] + ) + assert download_calls == [ + { + "model_name_or_path": "MiniMaxAI/MiniMax-H3", + "cache_dir": None, + "allow_patterns": expected_patterns, + "revision": None, + "require_all": True, + } + ] + + +def test_combined_weight_loader_routes_each_contiguous_partition(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + class FakeTransformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.loaded = [] + self.post_load_calls = 0 + + def load_weights(self, weights): + self.loaded = [name for name, _ in weights] + return set(self.loaded) + + def post_load_weights(self): + self.post_load_calls += 1 + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = FakeTransformer() + pipeline.transformers_ref = FakeTransformer() + pipeline.text_encoder = torch.nn.Identity() + pipeline.video_vae = torch.nn.Identity() + pipeline.audio_vae = torch.nn.Identity() + loaded = pipeline.load_weights( + iter( + [ + ("transformer.a", torch.ones(1)), + ("transformer.b", torch.ones(1)), + ("transformers_ref.a", torch.ones(1)), + ] + ) + ) + + assert pipeline.transformer.loaded == ["a", "b"] + assert pipeline.transformers_ref.loaded == ["a"] + assert pipeline.transformer.post_load_calls == 1 + assert pipeline.transformers_ref.post_load_calls == 1 + assert loaded == {"transformer.a", "transformer.b", "transformers_ref.a"} + + +def test_dlo_offload_plan_includes_token_refiner(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"} + assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"}) + + def test_joint_postprocess_is_multiprocessing_picklable(): from vllm_omni.diffusion.models.minimax_h3 import ( get_minimax_h3_post_process_func, @@ -87,6 +419,87 @@ def test_shifted_sigma_schedule_matches_reference_values(): ) +def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = True + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + # Model-side Ulysses shards rows before Attention gathers them again. + # The Python packed_total must therefore remain global even though q is + # local here. + q = torch.randn(4, 2, 4) + cu_seqlens = torch.tensor([0, 5, 8], dtype=torch.int32) + + output = attention._run_packed_attention( + q, + q, + q, + cu_seqlens=cu_seqlens, + max_seqlen=5, + packed_total=8, + ) + + assert output.shape == q.shape + assert attention.attention.metadata.attn_mask is None + assert attention.attention.metadata.extra["valid_kv_length"] == 5 + + +def test_packed_attention_keeps_padding_mask_for_other_backends(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = False + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + q = torch.randn(8, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + assert torch.equal( + attention.attention.metadata.attn_mask, + torch.tensor([[True, True, True, True, True, False, False, False]]), + ) + + def test_reference_image_resize_contract(): from PIL import Image @@ -293,6 +706,142 @@ def test_layerwise_offload_releases_text_encoder(): pipeline.text_encoder.offload_to_cpu.assert_called_once_with() +def test_distributed_layerwise_offload_releases_text_encoder(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + +def test_distributed_layerwise_offload_stages_vae_component(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + component = Mock() + + with pipeline._component_on_device(component): + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_not_called() + + component.offload_to_cpu.assert_called_once_with() + + +def test_distributed_layerwise_resident_blocks_are_stage_scoped(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + controller = Mock() + pipeline._dlo_residency_controller = controller + + with pipeline._resident_dit_layers_on_device(): + controller.load_resident_layers.assert_called_once_with() + controller.offload_resident_layers.assert_not_called() + + controller.offload_resident_layers.assert_called_once_with() + + +def test_distributed_layerwise_resident_blocks_release_on_failure(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + controller = Mock() + pipeline._dlo_residency_controller = controller + + with pytest.raises(RuntimeError, match="denoise failed"): + with pipeline._resident_dit_layers_on_device(): + raise RuntimeError("denoise failed") + + controller.load_resident_layers.assert_called_once_with() + controller.offload_resident_layers.assert_called_once_with() + + +def test_distributed_layerwise_resident_blocks_can_be_skipped(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + controller = Mock() + pipeline._dlo_residency_controller = controller + + with pipeline._resident_dit_layers_on_device(enabled=False): + pass + + controller.load_resident_layers.assert_not_called() + controller.offload_resident_layers.assert_not_called() + + +def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + ) + + class Stack(torch.nn.Module): + def __init__(self, count): + super().__init__() + self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) + + class TextStack(torch.nn.Module): + def __init__(self, count): + super().__init__() + self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) + + hooks = [] + + def fake_apply(block, next_block, device, stream, pin_memory): + hook = SimpleNamespace( + block=block, + next_block=next_block, + device=device, + stream=stream, + pin_memory=pin_memory, + _prev_hook=None, + offload_layer=Mock(), + ) + hooks.append(hook) + return hook + + monkeypatch.setattr( + "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook", + fake_apply, + ) + monkeypatch.setattr( + "vllm_omni.platforms.current_omni_platform.Stream", + Mock(return_value="copy-stream"), + ) + encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + encoder.device_target = torch.device("cpu") + encoder.vision = Stack(2) + encoder.text_model = TextStack(3) + + encoder.enable_omni_layerwise_offload(pin_memory=False) + + assert len(hooks) == 5 + assert all(hook._prev_hook is not None for hook in hooks) + assert all(hook.device == torch.device("cpu") for hook in hooks) + + def test_video_vae_keeps_reference_fp32_weights(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module @@ -324,6 +873,82 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): assert next(video_vae.parameters()).dtype == torch.float32 +def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1) + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("meta"), + load_device=torch.device("cpu"), + ) + + assert next(video_vae.parameters()).device.type == "cpu" + assert video_vae._device_target.type == "meta" + + +def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1) + + stager = Mock() + stager_cls = Mock(return_value=stager) + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("cuda"), + load_device=torch.device("cpu"), + ) + video_vae.load_to_device() + video_vae.offload_to_cpu() + + stager_cls.assert_called_once_with( + video_vae.remote, + torch.device("cuda"), + pin_memory=True, + ) + stager.load.assert_called_once_with() + stager.offload.assert_called_once_with() + + def test_video_vae_encode_uses_configured_parallel_tiling(): from vllm_omni.diffusion.models.minimax_h3.vae import ( MiniMaxH3VideoVAE, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index d6f089d0a..bcd074483 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py @@ -61,6 +61,33 @@ def test_packed_attention_is_a_regional_compile_boundary(): assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False) +def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims(): + from vllm_omni.diffusion.layers.rope import RotaryEmbedding + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + attention = object.__new__(MiniMaxH3Attention) + nn.Module.__init__(attention) + attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False) + attention.rope._forward_method = attention.rope.forward_native + + x = torch.randn(11, 3, 128, dtype=torch.bfloat16) + freqs_half = torch.randn(11, 48) + freqs = torch.cat((freqs_half, freqs_half), dim=-1) + actual = attention._apply_rope(x, freqs) + + cos = torch.cos(freqs).to(x.dtype).unsqueeze(1) + sin = torch.sin(freqs).to(x.dtype).unsqueeze(1) + x_rot = x[..., :96] + x1, x2 = x_rot.chunk(2, dim=-1) + expected_rot = x_rot * cos + torch.cat((-x2, x1), dim=-1) * sin + expected = torch.cat((expected_rot, x[..., 96:]), dim=-1) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0) + + @pytest.mark.parametrize( ("tp_size", "message"), [ diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index d93b1c9f8..80d99895f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -135,6 +135,7 @@ def _generate_joint_output(omni, config: QualityTestConfig): num_frames=config.num_frames, guidance_scale=config.guidance_scale, sigmas=config.sigmas, + extra_args={"aspect_ratio": "16:9"}, ), ) main_process_peak_mem = torch.accelerator.max_memory_allocated() / (1024**3) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index c02fdaabc..fd004d056 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -27,7 +27,9 @@ from vllm_omni.diffusion.offloader.block_discovery import ( from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( DistributedLayerwiseOffloadBackend, DistributedLayerwiseOffloadHook, + PinnedResidentLayerGroup, ) +from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan from vllm_omni.platforms import current_omni_platform @@ -97,6 +99,7 @@ def patched_offload_runtime(monkeypatch): monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent) monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None) class TinyBlock(nn.Module): @@ -308,6 +311,164 @@ class TestDistributedLayerwiseOffloadHook: assert shard1.numel() == 2 assert torch.equal(shard1, torch.tensor([102.0, 0.0])) + def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime): + block = nn.Module() + block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) + next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) + next_block.weight.mmap_weight_transform_pending = True + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(block) + + assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0])) + + def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime): + block = nn.Module() + block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) + next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(block) + + assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0])) + + +class TestPinnedResidentLayerGroup: + def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime): + blocks = [nn.Linear(2, 2), nn.Linear(2, 2)] + expected = [] + for index, block in enumerate(blocks): + block.weight.data.fill_(index + 1) + block.bias.data.fill_(index + 3) + expected.append((block.weight.detach().clone(), block.bias.detach().clone())) + + group = PinnedResidentLayerGroup( + blocks, + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + ) + + assert all(block.weight.numel() == 0 for block in blocks) + + group.load() + for block, (weight, bias) in zip(blocks, expected): + assert torch.equal(block.weight, weight) + assert torch.equal(block.bias, bias) + + # Device-side changes must not overwrite the persistent host master. + blocks[0].weight.data.zero_() + group.offload() + assert all(block.weight.numel() == 0 for block in blocks) + + group.load() + for block, (weight, bias) in zip(blocks, expected): + assert torch.equal(block.weight, weight) + assert torch.equal(block.bias, bias) + + group.offload() + + def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self): + module = nn.Linear(2, 2, bias=False) + + def transform(tensor): + return tensor.flip(0) + + module.weight.mmap_weight_transform = transform + + backend = object.__new__(DistributedLayerwiseOffloadBackend) + backend._remember_mmap_param_attrs(module) + module.to_empty(device="meta") + + replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float()) + backend._attach_mmap_param_attrs("weight", replacement, module.weight) + + assert replacement.mmap_weight_transform is transform + assert replacement.mmap_weight_transform_pending is True + + def test_shared_allgather_output_is_narrowed_to_current_block( + self, dist_group, patched_offload_runtime, monkeypatch + ): + block = nn.Module() + block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32)) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(block) + # Simulate unified buffers sized for a larger heterogeneous block. + hook.gpu_buffers = [ + {torch.float32: torch.empty(12)}, + {torch.float32: torch.empty(12)}, + ] + hook.gpu_shard_buffers = [ + {torch.float32: torch.empty(6)}, + {torch.float32: torch.empty(6)}, + ] + output_sizes = [] + + def fake_allgather(output, local_shard, *, group): + del group + output_sizes.append(output.numel()) + output.zero_() + output[: local_shard.numel()].copy_(local_shard) + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather) + + hook.prefetch_layer(slot=0, non_blocking=False) + + assert output_sizes == [4] + + +class TestPinnedModuleStager: + def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime): + module = nn.Linear(2, 2) + expected_weight = module.weight.detach().clone() + expected_bias = module.bias.detach().clone() + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + copy_stream=DummyStream(), + ) + + stager.load() + module.weight.data.zero_() + module.bias.data.zero_() + stager.offload() + + assert torch.equal(module.weight, expected_weight) + assert torch.equal(module.bias, expected_bias) + stager.load() + assert torch.equal(module.weight, expected_weight) + assert torch.equal(module.bias, expected_bias) + stager.offload() + class _DummyBlock(nn.Module): def __init__(self): @@ -504,6 +665,81 @@ class TestCrossGroupSharedBuffer: assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)" +class TestAllGatherOutputSlicing: + """Regression test: with heterogeneous block sizes (e.g. MiniMax H3's + token_refiner with blocks of ~1231MB and ~239MB), the shared AllGather + output buffers are sized to the *largest* block across all groups. + prefetch_layer must slice the output buffer down to each block's actual + AllGather output size, otherwise output.numel() != dp_size * input.numel() + and all_gather_into_tensor raises during enable()'s first prefetch. + """ + + def test_prefetch_slices_shared_output_buffer_to_block_size(self, dist_group, patched_offload_runtime, monkeypatch): + dp_size = 2 + + # Heterogeneous block sizes: 8 vs 32 elements. + small_block = nn.Module() + small_block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32)) + large_block = nn.Module() + large_block.weight = nn.Parameter(torch.arange(100, 132, dtype=torch.float32)) + + def make_hook(next_block: nn.Module) -> DistributedLayerwiseOffloadHook: + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=dist.group.WORLD, # non-None -> AllGather path (mocked below) + dp_size=dp_size, + rank=0, + pin_memory=False, + ) + current_block = nn.Module() + current_block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + hook.initialize_hook(current_block) + return hook + + hook_small = make_hook(small_block) + hook_large = make_hook(large_block) + hooks = [hook_small, hook_large] + + # Assign max-sized shared buffers exactly as backend.enable() does. + shared_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_buffers(hooks) + shared_shard_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_shard_buffers(hooks) + for hook in hooks: + hook.gpu_buffers = shared_buffers + hook.gpu_shard_buffers = shared_shard_buffers + + # Shared output buffer must be sized to the largest block (32). + assert shared_buffers[0][torch.float32].numel() == 32 + + # Fake AllGather that enforces the real torch size contract and + # emulates reconstruction of the full block from all ranks' shards. + full_values = { + 8: torch.arange(8, dtype=torch.float32), + 32: torch.arange(100, 132, dtype=torch.float32), + } + ag_output_numels: list[int] = [] + + def fake_all_gather(output, input, group=None): + assert output.numel() == dp_size * input.numel(), ( + f"AllGather contract violated: output.numel()={output.numel()} " + f"!= {dp_size} * input.numel()={input.numel()}" + ) + ag_output_numels.append(output.numel()) + output.copy_(full_values[output.numel()]) + + monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather) + + # Small block: output must be sliced to dp_size * shard = 8, not 32. + hook_small.prefetch_layer(slot=0, non_blocking=False) + assert ag_output_numels[-1] == 8 + assert torch.equal(small_block.weight, full_values[8]) + + # Large block: uses the full max-sized buffer. + hook_large.prefetch_layer(slot=1, non_blocking=False) + assert ag_output_numels[-1] == 32 + assert torch.equal(large_block.weight, full_values[32]) + + class TestOffloadPlan: """Test declarative OffloadPlan metadata.""" @@ -517,6 +753,7 @@ class TestOffloadPlan: plan = OffloadPlan( block_attrs={"transformer": ("gen_layers",)}, offload_submodules={"context_encoder": "layers"}, + resident_dit_paths=frozenset({"transformer"}), ) class PipelineWithPlan(nn.Module): @@ -527,12 +764,14 @@ class TestOffloadPlan: assert result is plan assert result.block_attrs == {"transformer": ("gen_layers",)} assert result.offload_submodules == {"context_encoder": "layers"} + assert result.resident_dit_paths == frozenset({"transformer"}) def test_offload_plan_defaults_to_empty(self): """OffloadPlan with no arguments should have empty dicts.""" plan = OffloadPlan() assert plan.block_attrs == {} assert plan.offload_submodules == {} + assert plan.resident_dit_paths == frozenset() def test_offload_plan_is_frozen(self): """OffloadPlan should be immutable (frozen=True).""" @@ -540,6 +779,50 @@ class TestOffloadPlan: with pytest.raises(Exception): plan.block_attrs = {"x": ("y",)} # type: ignore + def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime): + class TokenRefiner(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + self.token_refiner = TokenRefiner() + + class Pipeline(nn.Module): + _offload_plan = OffloadPlan( + offload_submodules={"token_refiner": "blocks"}, + resident_dit_paths=frozenset({"transformer"}), + ) + + def __init__(self): + super().__init__() + self.transformer = Transformer() + + pipeline = Pipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + dlo_resident_layers=2, + ), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert len(backend._resident_blocks) == 2 + assert len(backend._all_hook_groups) == 1 + assert len(backend._all_hook_groups[0]) == 2 + assert backend.enabled + + backend.disable() + class TestMmapValidation: """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" @@ -746,7 +1029,7 @@ class TestConfigValidation: OffloadConfig.from_od_config(FakeODConfig()) def test_hsdp_without_allgather_allowed(self): - """HSDP + DLO + no-AllGather should be allowed (full weights per rank).""" + """HSDP + DLO + no-AllGather uses standard-loader rank-local weights.""" from vllm_omni.diffusion.offloader.base import OffloadConfig class FakePC: @@ -768,12 +1051,33 @@ class TestConfigValidation: enable_layerwise_offload = False enable_distributed_layerwise_offload = True dlo_use_allgather = False # no AllGather → should be allowed + dlo_resident_layers = 20 pin_cpu_memory = True parallel_config = FakePC() model = "/fake/path" config = OffloadConfig.from_od_config(FakeODConfig()) assert config.dp_size == 1 # forced to 1 when no AllGather + assert config.dlo_resident_layers == 20 + + def test_resident_layers_with_allgather_rejected(self): + class FakePC: + data_parallel_size = 2 + use_hsdp = False + sequence_parallel_size = 1 + + class FakeODConfig: + enable_cpu_offload = False + enable_layerwise_offload = False + enable_distributed_layerwise_offload = True + dlo_use_allgather = True + dlo_resident_layers = 20 + pin_cpu_memory = True + parallel_config = FakePC() + model = "/fake/path" + + with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"): + OffloadConfig.from_od_config(FakeODConfig()) def test_num_inference_steps_none_rejected(self): """DP multi-concurrency should reject None num_inference_steps.""" diff --git a/tests/diffusion/quantization/test_int8_config.py b/tests/diffusion/quantization/test_int8_config.py index 875277ece..f706cc68b 100644 --- a/tests/diffusion/quantization/test_int8_config.py +++ b/tests/diffusion/quantization/test_int8_config.py @@ -147,6 +147,130 @@ def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyP assert isinstance(method, UnquantizedLinearMethod) +def test_fused_layer_can_be_ignored_by_its_prefix(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): + """A layer opts out of quantization under the prefix its module was built with. + + MiniMax H3 stores qkv as one checkpoint tensor, so the fused name is the + only handle a user has on it. + """ + from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + config = build_quant_config("int8", ignored_layers=["blocks.0.attn.qkv_proj"]) + layer = mocker.Mock(spec=LinearBase) + + assert isinstance(config.get_quant_method(layer, "blocks.0.attn.qkv_proj"), UnquantizedLinearMethod) + assert isinstance(config.get_quant_method(layer, "blocks.1.attn.qkv_proj"), NPUInt8OnlineLinearMethod) + assert isinstance(config.get_quant_method(layer, "blocks.0.mlp.fc1"), NPUInt8OnlineLinearMethod) + + +class TestNPUQuantMatmulShapeFallback: + """npu_quant_matmul rejects wide outputs, so those layers stay unquantized. + + The limit applies to the per-rank shard, which is what the kernel sees. The + same layer can therefore be quantized at a higher TP degree. + """ + + @pytest.fixture(autouse=True) + def _mock_tp(self, mocker): + # The fallback delegates to UnquantizedLinearMethod, which reads the TP + # group; stand one in so the shape gating can be tested without a + # distributed init. + mock_group = mocker.Mock() + mock_group.rank_in_group = 0 + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + + @pytest.fixture + def quant_config(self): + from vllm_omni.quantization.int8_config import DiffusionInt8Config + + return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic") + + def create_weights(self, method, layer, output_partition_sizes): + # Narrow input dim: the fallback path allocates for real, and only the + # output dimension is under test. + method.create_weights( + layer, + input_size_per_partition=8, + output_partition_sizes=output_partition_sizes, + input_size=8, + output_size=sum(output_partition_sizes), + params_dtype=torch.bfloat16, + weight_loader=lambda param, loaded_weight, *args, **kwargs: None, + ) + + @pytest.mark.parametrize( + "method_name", + ["NPUInt8LinearMethod", "NPUInt8OnlineLinearMethod"], + ) + def test_over_wide_layer_falls_back_to_unquantized(self, quant_config, method_name): + import vllm_omni.quantization.int8_config as int8_config + + method = getattr(int8_config, method_name)(quant_config) + layer = Module() + layer.quant_method = method + + self.create_weights(method, layer, [int8_config.NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1]) + + assert isinstance(layer.quant_method, UnquantizedLinearMethod) + assert layer.weight.dtype == torch.bfloat16 + + def test_layer_within_the_limit_keeps_int8(self, quant_config): + from vllm_omni.quantization.int8_config import ( + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + NPUInt8OnlineLinearMethod, + ) + + method = NPUInt8OnlineLinearMethod(quant_config) + layer = Module() + layer.quant_method = method + + self.create_weights(method, layer, [NPU_QUANT_MATMUL_MAX_OUT_FEATURES]) + + assert layer.quant_method is method + + def test_tensor_parallel_sharding_brings_a_wide_layer_back(self, quant_config): + from vllm_omni.quantization.int8_config import ( + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + NPUInt8OnlineLinearMethod, + ) + + global_out_features = 4 * NPU_QUANT_MATMUL_MAX_OUT_FEATURES + method = NPUInt8OnlineLinearMethod(quant_config) + layer = Module() + layer.quant_method = method + + self.create_weights(method, layer, [global_out_features // 4]) + + assert layer.quant_method is method + + +class TestOffloadAfterQuant: + def test_only_methods_advertising_the_capability_are_asked(self): + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + from vllm_omni.quantization.int8_config import DiffusionInt8Config, NPUInt8OnlineLinearMethod + + class MetaDeviceOnlyMethod: + """Stands in for upstream online FP8: lazy, but no offload hook.""" + + uses_meta_device = True + + model = torch.nn.Sequential(Module(), Module()) + lazy_int8 = NPUInt8OnlineLinearMethod( + DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic") + ) + model[0].quant_method = lazy_int8 + model[1].quant_method = MetaDeviceOnlyMethod() + + marked = DiffusersPipelineLoader._request_offload_after_quant(model) + + assert marked == 1 + assert lazy_int8._offload_after_quant + + class TestInt8LinearMethod: @pytest.fixture def mock_quant_config(self, mocker): diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 10fe00168..205b2e491 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -131,6 +131,11 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES +def test_task_type_roundtrip(): + od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task") + assert od.task_type == "model-defined-task" + + def test_additional_config_roundtrip(): additional_config = {"torchair_graph_config": {"enabled": True}} od = _roundtrip_diffusion_config(model="x", additional_config=additional_config) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 932630a23..f1c48ba30 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -247,6 +247,56 @@ class TestRequestModeDispatch: results = {ro.request_id: ro.result.error for ro in out.runner_outputs} assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids} + def test_dlo_dp_routes_multiple_requests_without_pipeline_batching(self): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + executor.execute_request = Mock(return_value="dlo-dp") + executor.collective_rpc = Mock() + scheduler_output = _make_sched_output("A", "B") + + result = executor.execute_batch(scheduler_output) + + assert result == "dlo-dp" + executor.execute_request.assert_called_once_with(scheduler_output) + executor.collective_rpc.assert_not_called() + + def test_dlo_dp_multi_rank_reply_uses_synchronous_rpc_collection(self): + executor, req_q, res_q = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + ) + executor._sync_result_buffer = res_q + + def worker(): + request = req_q.get(timeout=2) + assert "rpc_id" not in request + wave_id = request["wave_id"] + for dp_rank in range(2): + res_q.put( + { + "dp_rank": dp_rank, + "output": _tagged_output(str(dp_rank)), + "wave_id": wave_id, + } + ) + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + results = executor.collective_rpc( + "execute_model", + unique_reply_rank=None, + exec_all_ranks=True, + ) + thread.join(timeout=2) + + assert [result.error for result in results] == ["0", "1"] + # ───────────────── concurrent collective RPC ───────────────── diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index cf77e2256..ca491275e 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -128,6 +128,24 @@ def test_qwen3_tts_codec_frame_rate_patching(): assert omni_config.codec_frame_rate_hz == 12.3 +def test_qwen3_tts_startup_task_type_is_validated(): + vllm_config = EngineArgs().create_model_config() + + config = OmniModelConfig.from_vllm_model_config( + vllm_config, + model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM", + task_type="Base", + ) + assert config.task_type == "Base" + + with pytest.raises(ValueError, match="Qwen3-TTS --task-type must be one of"): + OmniModelConfig.from_vllm_model_config( + vllm_config, + model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM", + task_type="fl2va", + ) + + def test_from_cli_args_picks_up_stage_configs_path(): """from_cli_args should pick up stage_configs_path from namespace.""" ns = argparse.Namespace( diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py new file mode 100644 index 000000000..75660cbc7 --- /dev/null +++ b/tests/engine/test_stage_engine_args.py @@ -0,0 +1,510 @@ +from __future__ import annotations + +import copy +import sys +import types +from dataclasses import fields, replace +from pathlib import Path + +import pytest +from pydantic.fields import FieldInfo + +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, + PipelineConfig, + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, + StageType, + build_stage_runtime_overrides, + load_deploy_config, + merge_pipeline_deploy, +) +from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig +from vllm_omni.engine import stage_init_utils +from vllm_omni.engine.arg_utils import OmniEngineArgs +from vllm_omni.engine.stage_init_utils import ( + build_engine_args_dict, + build_engine_args_dict_from_omni_stage_config, + build_legacy_engine_args_dict, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy" + + +def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]: + backend_fields = fields(config_cls) + backend_field_names = {backend_field.name for backend_field in backend_fields} + backend_config = config_cls( + **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names} + ) + effective_values: dict[str, object] = {} + for backend_field in backend_fields: + value = getattr(backend_config, backend_field.name) + if isinstance(value, FieldInfo): + value = value.get_default(call_default_factory=True) + effective_values[backend_field.name] = value + return effective_values + + +_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs)) +_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig)) + + +@pytest.fixture(autouse=True) +def _stable_engine_arg_environment(monkeypatch, tmp_path): + from vllm_omni import platforms + + monkeypatch.delenv("VLLM_USE_FLASHINFER_MOE_FP16", raising=False) + xcodec_model = tmp_path / "xcodec-model" + xcodec_model.mkdir() + monkeypatch.setenv("XCODEC1_PATH", str(xcodec_model)) + platform = platforms.current_omni_platform + monkeypatch.setattr(platform, "device_name", "cpu", raising=False) + monkeypatch.setattr(platform, "device_type", "cpu", raising=False) + monkeypatch.setattr(platform, "is_rocm", lambda: False) + monkeypatch.setattr( + OmniDiffusionConfig, + "_resolve_master_port", + lambda config: config.master_port or 29500, + ) + + def _resolve_test_worker(engine_args): + worker_type = engine_args.get("worker_type") + if worker_type is not None and engine_args.get("worker_cls") in (None, "auto"): + engine_args["worker_cls"] = f"test.{worker_type}.Worker" + + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker) + + +def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]: + parent_model = tmp_path / "parent-model" + stage_model = tmp_path / "stage-model" + diffusion_model = tmp_path / "diffusion-model" + for root in (parent_model, stage_model, diffusion_model): + root.mkdir() + (stage_model / "ar-model").mkdir() + (stage_model / "ar-tokenizer").mkdir() + + pipeline = PipelineConfig( + model_type="typed-engine-args", + model_arch="PipelineModel", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + requires_multimodal_data=True, + hf_config_name="thinker_config", + engine_output_type="hidden_states", + model_arch="Qwen3OmniMoeForConditionalGeneration", + model_subdir="ar-model", + tokenizer_subdir="ar-tokenizer", + retains_state_across_chunks=True, + async_chunk_process_next_stage_input_func="operator.add", + ), + StagePipelineConfig( + stage_id=1, + model_stage="talker", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + engine_output_type="audio_tokens", + ), + StagePipelineConfig( + stage_id=2, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(1,), + final_output=True, + final_output_type="image", + model_arch="TypedDiffusionPipeline", + ), + ), + ) + deploy = DeployConfig( + async_chunk=True, + trust_remote_code=True, + dtype="float16", + distributed_executor_backend="mp", + stages=[ + StageDeployConfig( + stage_id=0, + tensor_parallel_size=2, + max_num_seqs=8, + default_sampling_params={"extra_args": {"voice": "test"}}, + engine_extras={ + "model": str(stage_model), + "attention_backend": "FLASHINFER", + "hf_overrides": {"rope_scaling": {"factor": 2}}, + "limit_mm_per_prompt": {"audio": 1}, + "logits_processors": ["test.custom.LogitsProcessor"], + "kv_cache_memory_bytes": 1024, + "omni_kv_config": {"need_send_cache": True}, + "max_cudagraph_capture_size": 0, + "worker_cls": "test.custom.Worker", + }, + ), + StageDeployConfig( + stage_id=1, + max_num_seqs=4, + ), + StageDeployConfig( + stage_id=2, + tensor_parallel_size=2, + vae_parallel_mode="spatial_shard_height", + diffusion_attention_config={ + "default": {"backend": "FLASH_ATTN"}, + "per_role": {"cross": {"backend": "TORCH_SDPA"}}, + }, + engine_extras={ + "model": str(diffusion_model), + "engine_backend": "test.diffusion.Engine", + "enable_session_state_manager": True, + }, + ), + ], + ) + return pipeline, deploy, str(parent_model) + + +def _legacy_and_typed_stages( + pipeline: PipelineConfig, + deploy: DeployConfig, + model: str, + cli_overrides: dict[str, object] | None = None, +): + resolved_cli_overrides = {"model": model, **(cli_overrides or {})} + legacy_stage_configs = merge_pipeline_deploy( + pipeline, + copy.deepcopy(deploy), + ) + for stage in legacy_stage_configs: + stage.runtime_overrides = build_stage_runtime_overrides( + stage.stage_id, + resolved_cli_overrides, + ) + legacy_stages = [stage.to_omegaconf() for stage in legacy_stage_configs] + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=copy.deepcopy(deploy), + cli_overrides=resolved_cli_overrides, + ) + return legacy_stages, omni_config + + +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) + connector_spec = {"name": "SharedMemoryConnector", "extra": {"mode": "test"}} + cli_tokenizer = "/external/tokenizer" + + typed_args_by_stage = {} + for stage_id in (0, 1): + legacy_args = build_legacy_engine_args_dict( + legacy_stages[stage_id], + model, + stage_connector_spec=connector_spec, + cli_tokenizer=cli_tokenizer, + ) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(stage_id), + model, + stage_connector_spec=connector_spec, + cli_tokenizer=cli_tokenizer, + ) + typed_args_by_stage[stage_id] = typed_args + + assert {name: typed_args[name] for name in legacy_args} == legacy_args + + thinker_args = typed_args_by_stage[0] + 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 + assert thinker_args["stage_connector_spec"] == connector_spec + assert thinker_args["has_sampling_extra_args"] is True + assert thinker_args["omni_kv_config"] == {"need_send_cache": True} + assert thinker_args["worker_cls"] == "test.custom.Worker" + assert thinker_args["attention_backend"] == "FLASHINFER" + assert thinker_args["trust_remote_code"] is True + assert thinker_args["dtype"] == "float16" + assert thinker_args["distributed_executor_backend"] == "mp" + assert thinker_args["tensor_parallel_size"] == 2 + assert thinker_args["hf_overrides"] == {"rope_scaling": {"factor": 2}} + assert thinker_args["limit_mm_per_prompt"] == {"audio": 1} + assert thinker_args["logits_processors"] == ["test.custom.LogitsProcessor"] + assert thinker_args["retains_state_across_chunks"] is True + assert thinker_args["kv_cache_memory_bytes"] == 1024 + assert thinker_args["max_cudagraph_capture_size"] == 0 + + talker_args = typed_args_by_stage[1] + assert talker_args["model"] == model + assert talker_args["tokenizer"] == cli_tokenizer + assert talker_args["worker_cls"] == "test.generation.Worker" + assert talker_args["disable_hybrid_kv_cache_manager"] is True + assert talker_args["enable_prefix_caching"] is False + + +def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(2), + model, + ) + + assert typed_args["stage_id"] == legacy_args["stage_id"] == 2 + assert typed_args["model"] == legacy_args["model"] == str(tmp_path / "diffusion-model") + assert omni_config.stage_by_id(2).diffusion_config.model == str(tmp_path / "diffusion-model") + assert typed_args["model_stage"] == legacy_args["model_stage"] == "dit" + assert typed_args["model_arch"] == legacy_args["model_arch"] == "TypedDiffusionPipeline" + assert typed_args["model_class_name"] == legacy_args["model_class_name"] == "TypedDiffusionPipeline" + assert typed_args["engine_backend"] == legacy_args["engine_backend"] == "test.diffusion.Engine" + assert typed_args["enable_session_state_manager"] is legacy_args["enable_session_state_manager"] is True + assert typed_args["parallel_config"]["tensor_parallel_size"] == 2 + assert typed_args["parallel_config"]["vae_parallel_mode"] == "spatial_shard_height" + assert isinstance(typed_args["diffusion_attention_config"], AttentionConfig) + assert typed_args["diffusion_attention_config"].default.backend == "FLASH_ATTN" + assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA" + + +def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + deploy.enable_prefix_caching = False + deploy.enable_chunked_prefill = False + thinker_deploy = deploy.stages[0] + thinker_deploy.gpu_memory_utilization = 0.75 + thinker_deploy.disable_hybrid_kv_cache_manager = False + thinker_deploy.async_scheduling = False + legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(0), + model, + ) + expected = { + "gpu_memory_utilization": 0.75, + "enable_prefix_caching": False, + "disable_hybrid_kv_cache_manager": False, + "max_num_seqs": 8, + "enable_chunked_prefill": False, + "async_scheduling": False, + } + + assert {name: legacy_args[name] for name in expected} == expected + assert {name: typed_args[name] for name in expected} == expected + + +@pytest.mark.parametrize("stage_scoped", [False, True], ids=["global", "stage-scoped"]) +def test_typed_engine_args_preserve_inherited_model_and_load_cli_fields(tmp_path, stage_scoped): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + expected = { + "revision": "model-rev", + "tokenizer_revision": "tokenizer-rev", + "code_revision": "code-rev", + "seed": 123, + "download_dir": str(tmp_path / "downloads"), + } + cli_overrides = {f"stage_0_{name}" if stage_scoped else name: value for name, value in expected.items()} + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model, + cli_overrides, + ) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(0), + model, + ) + + assert {name: legacy_args[name] for name in expected} == expected + assert {name: typed_args[name] for name in expected} == expected + + +def test_typed_engine_args_preserve_deploy_subdirectory_precedence(tmp_path): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + stage_model = tmp_path / "stage-model" + (stage_model / "override-model").mkdir() + (stage_model / "override-tokenizer").mkdir() + deploy.stages[0].engine_extras.update( + { + "model_subdir": "override-model", + "tokenizer_subdir": "override-tokenizer", + } + ) + legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) + + typed_stage = omni_config.stage_by_id(0) + legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model) + typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model) + + assert typed_stage.model_config.model_subdir == "override-model" + assert typed_stage.model_config.tokenizer_subdir == "override-tokenizer" + assert typed_args["model"] == legacy_args["model"] == str(stage_model / "override-model") + assert typed_args["tokenizer"] == legacy_args["tokenizer"] == str(stage_model / "override-tokenizer") + + +@pytest.mark.parametrize( + ("cli_overrides", "expected"), + [ + ({}, {"need_recv_cache": True}), + ( + {"stage_0_omni_kv_config": {"need_send_cache": False}}, + {"need_send_cache": False}, + ), + ], + ids=["topology-over-deploy", "cli-over-topology"], +) +def test_typed_engine_args_preserve_legacy_omni_kv_precedence(tmp_path, cli_overrides, expected): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + thinker = replace( + pipeline.stages[0], + omni_kv_config={"need_recv_cache": True}, + ) + pipeline = replace(pipeline, stages=(thinker, *pipeline.stages[1:])) + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model, + cli_overrides, + ) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model) + typed_stage = omni_config.stage_by_id(0) + typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model) + + assert typed_stage.connector_config.omni_kv_config == expected + assert typed_args["omni_kv_config"] == legacy_args["omni_kv_config"] == expected + + +def test_build_engine_args_dict_preserves_legacy_api(tmp_path): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + legacy_stages, _ = _legacy_and_typed_stages(pipeline, deploy, model) + connector_spec = {"name": "SharedMemoryConnector", "extra": {}} + + compatibility_args = build_engine_args_dict( + copy.deepcopy(legacy_stages[0]), + model, + stage_connector_spec=connector_spec, + cli_tokenizer="/external/tokenizer", + ) + explicitly_legacy_args = build_legacy_engine_args_dict( + copy.deepcopy(legacy_stages[0]), + model, + stage_connector_spec=connector_spec, + cli_tokenizer="/external/tokenizer", + ) + + assert compatibility_args == explicitly_legacy_args + + +def test_legacy_engine_args_drop_none_tp_without_mutating_stage_config(): + stage_config = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"tensor_parallel_size": None}, + default_sampling_params={}, + ) + + engine_args = build_legacy_engine_args_dict( + stage_config, + model="test-model", + ) + + assert "tensor_parallel_size" not in engine_args + assert stage_config.engine_args == {"tensor_parallel_size": None} + + +def test_typed_engine_args_own_rocm_attention_default(monkeypatch, tmp_path): + from vllm_omni import platforms + + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) + monkeypatch.setattr(platforms.current_omni_platform, "is_rocm", lambda: True) + aiter_module = types.ModuleType("vllm._aiter_ops") + setattr( + aiter_module, + "rocm_aiter_ops", + types.SimpleNamespace(is_enabled=lambda: False), + ) + monkeypatch.setitem(sys.modules, "vllm._aiter_ops", aiter_module) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model) + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(1), + model, + ) + + assert "attention_backend" not in legacy_args + assert typed_args["attention_backend"] == "TRITON_ATTN" + + +def test_typed_ming_image_engine_args_defer_diffusion_batch_default(): + pipeline = resolve_pipeline_config("ming_flash_omni_image") + assert pipeline is not None + deploy = load_deploy_config(_DEPLOY_DIR / "ming_flash_omni_image.yaml") + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model="/tmp", + ) + + legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model="/tmp") + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(1), + model="/tmp", + ) + typed_backend_args = {name: value for name, value in typed_args.items() if name in _DIFFUSION_BACKEND_FIELDS} + + assert "max_num_seqs" not in legacy_args + assert "max_num_seqs" not in typed_args + assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1 + + +@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) +def test_typed_engine_args_match_current_registry_backend_semantics(model_type): + pipeline = resolve_pipeline_config(model_type) + if pipeline is None: + pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve") + + deploy = ( + load_deploy_config(_DEPLOY_DIR / pipeline.default_deploy_config_name) + if pipeline.default_deploy_config_name is not None + else DeployConfig() + ) + legacy_stages, omni_config = _legacy_and_typed_stages( + pipeline, + deploy, + model="/tmp", + ) + + for legacy_stage in legacy_stages: + stage_id = legacy_stage.stage_id + legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp") + typed_args = build_engine_args_dict_from_omni_stage_config( + omni_config.stage_by_id(stage_id), + model="/tmp", + ) + if legacy_stage.stage_type == StageType.DIFFUSION: + backend_fields = _DIFFUSION_BACKEND_FIELDS + backend_config_cls = OmniDiffusionConfig + else: + backend_fields = _LLM_BACKEND_FIELDS + backend_config_cls = OmniEngineArgs + + missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys() + assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}" + + legacy_effective_args = _effective_backend_values(backend_config_cls, legacy_args) + typed_effective_args = _effective_backend_values(backend_config_cls, typed_args) + assert typed_effective_args == legacy_effective_args, ( + f"{model_type} stage {stage_id} changed effective backend arguments" + ) diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 652adef4d..08d223c68 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -48,6 +48,52 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): assert mux_calls[0]["audio"] is None +def test_float_frames_are_converted_without_stacking_full_video(monkeypatch): + frame = np.array( + [ + [[0.0, 0.5, 1.0, 0.2], [1.5, -0.5, 0.5, 0.8]], + ], + dtype=np.float32, + ) + original = frame.copy() + + def fail_stack(*args, **kwargs): + raise AssertionError("float video conversion must not stack all frames") + + monkeypatch.setattr(video_api_utils.np, "stack", fail_stack) + + frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame]) + + expected = np.array([[[128, 191, 255], [255, 64, 191]]], dtype=np.uint8) + assert frames.flags.c_contiguous + np.testing.assert_array_equal(frames, np.array([expected, expected])) + np.testing.assert_array_equal(frame, original) + + +def test_two_dimensional_width_four_frames_are_not_treated_as_rgba(): + frame = np.arange(8, dtype=np.float32).reshape(2, 4) / 8 + + frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame]) + + expected = np.rint(frame * 255).astype(np.uint8) + assert frames.shape == (2, 2, 4) + np.testing.assert_array_equal(frames, np.stack([expected, expected])) + + +def test_mixed_float_dtypes_preserve_stacked_rounding_semantics(): + frames = [ + np.full((1, 1, 3), 0.1, dtype=np.float16), + np.full((1, 1, 3), 0.2, dtype=np.float32), + ] + stacked = np.stack(frames) + expected = np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8) + + actual = video_api_utils._coerce_video_to_uint8_frames(frames) + + assert expected[0, 0, 0, 0] == 25 + np.testing.assert_array_equal(actual, expected) + + @pytest.mark.parametrize("frame_count", [3, 4]) def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count): video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 8987ae1d8..ecb0f20e8 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -240,6 +240,28 @@ def test_serve_cli_accepts_ulysses_mode(): assert parallel_config.ulysses_mode == "advanced_uaa" +def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage(): + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--task-type", + "fl2va", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + + assert args.task_type == "fl2va" + assert stage_cfg["engine_args"]["task_type"] == "fl2va" + + def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): """Ensure diffusion serve CLI exposes the profiler switch.""" parser = TrackingArgumentParser() @@ -262,6 +284,36 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True +def test_serve_cli_forwards_distributed_offload_residency(): + """Ensure the two-GPU DLO placement controls reach the diffusion stage.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--enable-distributed-layerwise-offload", + "--dlo-no-use-allgather", + "--dlo-resident-layers", + "20", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + engine_args = stage_cfg["engine_args"] + + assert args.enable_distributed_layerwise_offload is True + assert args.dlo_use_allgather is False + assert args.dlo_resident_layers == 20 + assert engine_args["enable_distributed_layerwise_offload"] is True + assert engine_args["dlo_use_allgather"] is False + assert engine_args["dlo_resident_layers"] == 20 + + def test_serve_cli_accepts_diffusion_compile_controls(): """Ensure both compile controls reach the diffusion stage.""" parser = TrackingArgumentParser() diff --git a/tests/entrypoints/test_omni_snapshot_download.py b/tests/entrypoints/test_omni_snapshot_download.py index 412c0f6d9..d2c76a249 100644 --- a/tests/entrypoints/test_omni_snapshot_download.py +++ b/tests/entrypoints/test_omni_snapshot_download.py @@ -53,6 +53,11 @@ def download_backend(monkeypatch: pytest.MonkeyPatch): "download_weights_from_hf_specific", lambda **_kwargs: picked.append("huggingface"), ) + monkeypatch.setattr( + omni_base, + "file_or_path_exists", + lambda *_args, **_kwargs: False, + ) def run() -> str: picked.clear() @@ -82,3 +87,14 @@ def test_unset_defaults_to_huggingface(monkeypatch, download_backend): monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) assert download_backend() == "huggingface" + + +def test_modular_diffusers_defers_component_download(monkeypatch, download_backend): + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + monkeypatch.setattr( + omni_base, + "file_or_path_exists", + lambda *_args, **_kwargs: True, + ) + + assert download_backend() == "none" diff --git a/tests/helpers/mark.py b/tests/helpers/mark.py index 3475a0012..676732a8e 100644 --- a/tests/helpers/mark.py +++ b/tests/helpers/mark.py @@ -2,10 +2,69 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …).""" +from __future__ import annotations + +import re +from functools import lru_cache +from pathlib import Path + import pytest from vllm.platforms import current_platform -# Re-exported from tests.helpers.env (GPU wait + DeviceMemoryMonitor). +# Marker description tag in ``pyproject.toml`` ``tool.pytest.ini_options.markers``. +# Example: ``"H100: [hardware-resource] Tests that require H100 GPU"``. +_HARDWARE_RESOURCE_MARKER_TAG = "[hardware-resource]" + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[2] + + +def _load_pyproject() -> dict: + path = _repo_root() / "pyproject.toml" + try: + try: + import tomllib + except ModuleNotFoundError: # Python < 3.11 + import tomli as tomllib # type: ignore[no-redef] + + with path.open("rb") as f: + return tomllib.load(f) + except Exception: + return {} + + +@lru_cache(maxsize=1) +def get_hardware_mark_list() -> frozenset[str]: + """Return hardware SKU marker names from ``pyproject.toml``. + + A marker is treated as a hardware resource marker when its description + contains the ``[hardware-resource]`` tag, e.g.:: + + "H100: [hardware-resource] Tests that require H100 GPU" + """ + data = _load_pyproject() + entries = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("markers", []) + if not entries: + # Fallback when TOML parsing fails: scan tagged marker lines. + text = (_repo_root() / "pyproject.toml").read_text(encoding="utf-8") + return frozenset( + re.findall( + rf'^\s*"([A-Za-z0-9_]+)\s*:\s*{re.escape(_HARDWARE_RESOURCE_MARKER_TAG)}', + text, + flags=re.M, + ) + ) + + names: set[str] = set() + for entry in entries: + text = str(entry) + if _HARDWARE_RESOURCE_MARKER_TAG not in text: + continue + name = text.split(":", 1)[0].strip() + if name: + names.add(name) + return frozenset(names) def cuda_marks(*, res: str, num_cards: int): diff --git a/tests/model_executor/models/minicpmo_4_5/test_audio_chunk_mask.py b/tests/model_executor/models/minicpmo_4_5/test_audio_chunk_mask.py new file mode 100644 index 000000000..7769b38be --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_audio_chunk_mask.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniCPM-o 4.5 audio chunk attention masks.""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _reference_chunk_mask( + size: int, + chunk_size: int, + num_left_chunks: int, + num_lookhead: int, +) -> torch.Tensor: + """Preserve the original row-by-row implementation as a test oracle.""" + mask = torch.zeros(size, size, dtype=torch.bool) + for row in range(size): + if num_left_chunks < 0: + start = 0 + else: + start = max( + (row // chunk_size - num_left_chunks) * chunk_size, + 0, + ) + end = min( + (row // chunk_size + 1) * chunk_size + num_lookhead, + size, + ) + mask[row, start:end] = True + return mask + + +@pytest.mark.parametrize("size", [0, 1, 49, 50, 51, 1500]) +@pytest.mark.parametrize("chunk_size", [1, 17, 50]) +@pytest.mark.parametrize("num_left_chunks", [-1, 0, 1, 3]) +@pytest.mark.parametrize("num_lookhead", [0, 1, 7]) +def test_subsequent_chunk_mask_matches_row_by_row_reference( + size: int, + chunk_size: int, + num_left_chunks: int, + num_lookhead: int, +) -> None: + actual = MiniCPMO45OmniLLMForConditionalGeneration.subsequent_chunk_mask( + None, + size=size, + chunk_size=chunk_size, + num_left_chunks=num_left_chunks, + num_lookhead=num_lookhead, + ) + expected = _reference_chunk_mask( + size, + chunk_size, + num_left_chunks, + num_lookhead, + ) + + assert actual.dtype == torch.bool + assert actual.shape == (size, size) + assert torch.equal(actual, expected) diff --git a/tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py b/tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py new file mode 100644 index 000000000..3139dbeac --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression coverage for MiniCPM-o 4.5 SigLIP FlashAttention.""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch +from pytest_mock import MockerFixture +from transformers.utils import is_flash_attn_2_available + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni_llm as minicpmo_model +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + SiglipAttention, + SiglipEncoder, + SiglipFlashAttention2, + SiglipVisionConfig, +) + +_VISION_NUM_HEADS = 16 +_VISION_HEAD_DIM = 72 +_VISION_HIDDEN_SIZE = _VISION_NUM_HEADS * _VISION_HEAD_DIM +_DTYPE = torch.bfloat16 + + +def _attention_config(*, hidden_size: int, num_heads: int) -> SiglipVisionConfig: + return SiglipVisionConfig( + hidden_size=hidden_size, + num_attention_heads=num_heads, + attention_dropout=0.0, + ) + + +def _encoder_config(*, num_layers: int = 3) -> SiglipVisionConfig: + config = SiglipVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=num_layers, + num_attention_heads=4, + attention_dropout=0.0, + ) + config._attn_implementation = "flash_attention_2" + return config + + +def _capture_output( + outputs: dict[str, torch.Tensor], name: str +) -> Callable[[torch.nn.Module, tuple[torch.Tensor, ...], torch.Tensor], None]: + def hook( + _module: torch.nn.Module, + _inputs: tuple[torch.Tensor, ...], + output: torch.Tensor, + ) -> None: + outputs[name] = output + + return hook + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_qkv_bshd_layout_reuses_projection_storage(mocker: MockerFixture) -> None: + batch_size = 2 + seq_len = 7 + hidden_size = 32 + num_heads = 4 + head_dim = hidden_size // num_heads + attention = SiglipFlashAttention2(_attention_config(hidden_size=hidden_size, num_heads=num_heads)).eval() + + projected: dict[str, torch.Tensor] = {} + attention.q_proj.register_forward_hook(_capture_output(projected, "query")) + attention.k_proj.register_forward_hook(_capture_output(projected, "key")) + attention.v_proj.register_forward_hook(_capture_output(projected, "value")) + + flash_inputs: dict[str, torch.Tensor] = {} + + def fake_flash_attention( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + _attention_mask: torch.Tensor | None, + _query_length: int, + *, + dropout: float = 0.0, + softmax_scale: float | None = None, + unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None, + ) -> torch.Tensor: + del dropout, softmax_scale, unpadding_metadata + assert _attention_mask is None + assert _query_length == seq_len + flash_inputs.update( + query=query_states, + key=key_states, + value=value_states, + ) + return query_states + + flash_forward = mocker.patch.object( + attention, + "_flash_attention_forward", + side_effect=fake_flash_attention, + ) + + hidden_states = torch.randn(batch_size, seq_len, hidden_size) + with torch.inference_mode(): + output, _ = attention(hidden_states) + + assert output.shape == hidden_states.shape + flash_forward.assert_called_once() + + expected_shape = (batch_size, seq_len, num_heads, head_dim) + expected_stride = ( + seq_len * num_heads * head_dim, + num_heads * head_dim, + head_dim, + 1, + ) + for name in ("query", "key", "value"): + projected_tensor = projected[name] + flash_tensor = flash_inputs[name] + assert flash_tensor.shape == expected_shape + assert flash_tensor.stride() == expected_stride + assert flash_tensor.is_contiguous() + assert flash_tensor.storage_offset() == projected_tensor.storage_offset() + assert flash_tensor.untyped_storage().data_ptr() == projected_tensor.untyped_storage().data_ptr() + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("attention_mask", "expected_calls"), + [ + pytest.param(torch.tensor([[True, True, False], [True, False, False]]), 1, id="padded"), + pytest.param(None, 0, id="dense"), + ], +) +def test_encoder_computes_unpadding_metadata_once_per_forward( + mocker: MockerFixture, + attention_mask: torch.Tensor | None, + expected_calls: int, +) -> None: + encoder = SiglipEncoder(_encoder_config()).eval() + metadata_spy = mocker.spy(minicpmo_model, "_get_unpad_data") + received_metadata: list[minicpmo_model.FlashAttentionUnpaddingMetadata | None] = [] + + def fake_flash_attention( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + _attention_mask: torch.Tensor | None, + _query_length: int, + *, + dropout: float = 0.0, + softmax_scale: float | None = None, + unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None, + ) -> torch.Tensor: + del key_states, value_states, _attention_mask, _query_length, dropout, softmax_scale + received_metadata.append(unpadding_metadata) + return query_states + + for layer in encoder.layers: + mocker.patch.object(layer.self_attn, "_flash_attention_forward", side_effect=fake_flash_attention) + + hidden_states = torch.randn(2, 3, 32) + with torch.inference_mode(): + encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + ) + + assert metadata_spy.call_count == expected_calls + assert len(received_metadata) == len(encoder.layers) + if attention_mask is None: + assert all(metadata is None for metadata in received_metadata) + else: + assert received_metadata[0] is not None + assert all(metadata is received_metadata[0] for metadata in received_metadata) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_encoder_recomputes_unpadding_metadata_for_new_mask(mocker: MockerFixture) -> None: + encoder = SiglipEncoder(_encoder_config(num_layers=1)).eval() + metadata_spy = mocker.spy(minicpmo_model, "_get_unpad_data") + received_metadata: list[minicpmo_model.FlashAttentionUnpaddingMetadata | None] = [] + + def fake_flash_attention( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + _attention_mask: torch.Tensor | None, + _query_length: int, + *, + dropout: float = 0.0, + softmax_scale: float | None = None, + unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None, + ) -> torch.Tensor: + del key_states, value_states, _attention_mask, _query_length, dropout, softmax_scale + received_metadata.append(unpadding_metadata) + return query_states + + mocker.patch.object( + encoder.layers[0].self_attn, + "_flash_attention_forward", + side_effect=fake_flash_attention, + ) + hidden_states = torch.randn(2, 4, 32) + masks = ( + torch.tensor([[True, True, True, False], [True, True, False, False]]), + torch.tensor([[True, True, False, False], [True, False, False, False]]), + ) + + with torch.inference_mode(): + for attention_mask in masks: + encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + ) + + assert metadata_spy.call_count == len(masks) + first_metadata, second_metadata = received_metadata + assert first_metadata is not None + assert second_metadata is not None + assert first_metadata is not second_metadata + assert not torch.equal(first_metadata[0], second_metadata[0]) + + +def _require_flash_attention() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for FlashAttention execution") + if not is_flash_attn_2_available(): + pytest.skip("Transformers reports that FlashAttention 2 is unavailable") + + +def _make_cuda_attention_pair() -> tuple[SiglipAttention, SiglipFlashAttention2]: + config = _attention_config( + hidden_size=_VISION_HIDDEN_SIZE, + num_heads=_VISION_NUM_HEADS, + ) + torch.manual_seed(0) + eager_attention = SiglipAttention(config).eval().to(device="cuda", dtype=_DTYPE) + flash_attention = SiglipFlashAttention2(config).eval().to(device="cuda", dtype=_DTYPE) + flash_attention.load_state_dict(eager_attention.state_dict()) + return eager_attention, flash_attention + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}) +def test_encoder_metadata_reuse_matches_per_layer_recomputation() -> None: + _require_flash_attention() + encoder = SiglipEncoder(_encoder_config()).eval().to(device="cuda", dtype=_DTYPE) + torch.manual_seed(3) + hidden_states = torch.randn(2, 64, 32, device="cuda", dtype=_DTYPE) + attention_mask = torch.arange(64, device="cuda").unsqueeze(0) < torch.tensor( + (64, 43), + device="cuda", + ).unsqueeze(1) + + with torch.inference_mode(): + expected = hidden_states + for encoder_layer in encoder.layers: + expected = encoder_layer( + expected, + attention_mask, + output_attentions=False, + unpadding_metadata=None, + )[0] + actual = encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + )[0] + + torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}) +def test_flash_attention_unpadded_matches_eager_reference() -> None: + _require_flash_attention() + eager_attention, flash_attention = _make_cuda_attention_pair() + torch.manual_seed(1) + hidden_states = torch.randn( + 2, + 64, + _VISION_HIDDEN_SIZE, + device="cuda", + dtype=_DTYPE, + ) + + with torch.inference_mode(): + expected, _ = eager_attention(hidden_states) + actual, _ = flash_attention(hidden_states) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}) +@pytest.mark.parametrize( + "valid_lengths", + [ + pytest.param((64, 43), id="moderate_padding"), + pytest.param((64, 17), id="high_padding"), + ], +) +def test_flash_attention_padded_matches_per_sample_unpadded( + valid_lengths: tuple[int, int], +) -> None: + _require_flash_attention() + _, flash_attention = _make_cuda_attention_pair() + batch_size = len(valid_lengths) + seq_len = max(valid_lengths) + torch.manual_seed(2) + hidden_states = torch.randn( + batch_size, + seq_len, + _VISION_HIDDEN_SIZE, + device="cuda", + dtype=_DTYPE, + ) + attention_mask = torch.arange(seq_len, device="cuda").unsqueeze(0) < torch.tensor( + valid_lengths, + device="cuda", + ).unsqueeze(1) + + with torch.inference_mode(): + padded_output, _ = flash_attention( + hidden_states, + attention_mask=attention_mask, + ) + unpadded_outputs = [ + flash_attention(hidden_states[index : index + 1, :valid_length])[0] + for index, valid_length in enumerate(valid_lengths) + ] + + for index, valid_length in enumerate(valid_lengths): + torch.testing.assert_close( + padded_output[index : index + 1, :valid_length], + unpadded_outputs[index], + atol=2e-2, + rtol=2e-2, + ) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index e93abb74f..f5d2768c5 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -33,6 +33,7 @@ EXCLUDED_MODELS = [ "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", "MiniMaxH3Pipeline", + "MiniMaxH3ModularPipeline", "StableAudioPipeline", "WanImageToVideoPipeline", "WanS2VPipeline", diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index d172347f7..d18cabcf5 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -17,6 +17,8 @@ import vllm_omni.model_executor.models as me_models logger = init_logger(__name__) +_QWEN3_TTS_TASK_TYPES = frozenset({"CustomVoice", "VoiceDesign", "Base"}) + class OmniModelArchConfigConvertor(ModelArchConfigConvertorBase): """Config convertor for Omni multi-stage models. @@ -101,8 +103,8 @@ class OmniModelConfig(ModelConfig): "audio", "latents"). If None, output type is inferred. stage_connector_config: Stage connector configuration dictionary. Contains "name" (connector name), "extra" (extra connector config). - task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base). - If not specified, will be inferred from model path. + task_type: Model-defined startup task type. Each model validates its + supported values and applies the corresponding behavior. The correct way to initialize this class is via vLLM config, as most @@ -208,6 +210,14 @@ class OmniModelConfig(ModelConfig): ) return get_hf_text_config(self.hf_config) + def _validate_startup_task_type(self) -> None: + """Validate startup-only task selectors owned by an AR model.""" + if self.model_arch != "Qwen3TTSTalkerForConditionalGenerationARVLLM" or self.task_type is None: + return + if self.task_type not in _QWEN3_TTS_TASK_TYPES: + supported = ", ".join(sorted(_QWEN3_TTS_TASK_TYPES)) + raise ValueError(f"Qwen3-TTS --task-type must be one of {supported}; got {self.task_type!r}") + def _patch_qwen3_tts(self): """Patches the value of `position_id_per_seconds` in Qwen3's TTS's talker_config into the this class's codec_frame_rate_hz. @@ -275,6 +285,7 @@ class OmniModelConfig(ModelConfig): omni_cfg._patch_qwen3_tts() omni_cfg._maybe_override_text_config() + omni_cfg._validate_startup_task_type() if omni_cfg.hf_config is not None: omni_cfg.hf_config.architectures = omni_cfg.architectures diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 0f523eb48..0ddcb2b37 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -46,6 +46,29 @@ _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | _PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS +_NON_STAGE_ENGINE_CLI_FIELDS = frozenset( + { + "async_chunk", + "model", + "omni", + "output_modalities", + "stage_configs_path", + "stage_id", + "tokenizer", + } +) + +# Legacy deploy files can still carry StageConfig metadata in engine_extras. +# Those values are not backend engine inputs, and the additive typed path keeps +# sourcing their effective values from the immutable pipeline topology. +_LEGACY_STAGE_METADATA_EXTRA_FIELDS = frozenset( + { + "final_output", + "final_output_type", + "is_comprehension", + } +) + _QuantizationConfigType: TypeAlias = QuantizationConfig | str | Mapping[str, Any] | None @@ -55,6 +78,21 @@ class _QuantizationEngineOverrides(TypedDict, total=False): class _ModelEngineOverrides(TypedDict, total=False): + model: str + model_arch: str + model_subdir: str + tokenizer_subdir: str + revision: str + tokenizer_revision: str + code_revision: str + seed: int + logits_processors: list[str | type] + trust_remote_code: bool + dtype: Any + attention_backend: Any + moe_backend: str + hf_overrides: Any + limit_mm_per_prompt: dict[str, Any] active_stream_window: int enable_sleep_mode: bool subtalker_sampling_params: dict[str, Any] @@ -63,6 +101,7 @@ class _ModelEngineOverrides(TypedDict, total=False): task_type: str codec_frame_rate_hz: float enforce_eager: bool + max_cudagraph_capture_size: int enable_flashinfer_autotune: bool compilation_config: dict[str, Any] enable_multithread_weight_load: bool @@ -71,6 +110,9 @@ class _ModelEngineOverrides(TypedDict, total=False): class _LoadEngineOverrides(TypedDict, total=False): + tokenizer: str + download_dir: str + skip_tokenizer_init: bool load_format: str tokenizer_mode: str config_format: str @@ -78,6 +120,7 @@ class _LoadEngineOverrides(TypedDict, total=False): class _CacheEngineOverrides(TypedDict, total=False): + kv_cache_memory_bytes: int gpu_memory_utilization: float enable_prefix_caching: bool disable_hybrid_kv_cache_manager: bool @@ -93,6 +136,8 @@ class _SchedulerEngineOverrides(TypedDict, total=False): class _RuntimeEngineOverrides(TypedDict, total=False): + distributed_executor_backend: Any + worker_cls: str devices: str num_replicas: int env: dict[str, Any] @@ -113,6 +158,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False): ulysses_mode: str cfg_parallel_size: int vae_patch_parallel_size: int + vae_parallel_mode: str text_encoder_tp_size: int use_hsdp: bool mask_sp_padding: bool @@ -125,6 +171,10 @@ class _ParallelEngineOverrides(_ParallelConfigEngineOverrides, total=False): parallel_config: _ParallelConfigEngineOverrides | Mapping[str, Any] +class _ConnectorEngineOverrides(TypedDict, total=False): + omni_kv_config: dict[str, Any] + + @dataclass(frozen=True) class _StageEngineValues: """Typed projections of legacy flat per-stage ``yaml_engine_args``.""" @@ -134,6 +184,7 @@ class _StageEngineValues: load: _LoadEngineOverrides cache: _CacheEngineOverrides scheduler: _SchedulerEngineOverrides + connector: _ConnectorEngineOverrides runtime: _RuntimeEngineOverrides parallel: _ParallelEngineOverrides diffusion: _DiffusionEngineOverrides @@ -255,8 +306,21 @@ def _get_deploy_config( @config class OmniStageModelConfig: - """Per-stage model behavior.""" + """Per-stage model behavior and resolved model-engine inputs.""" + model: str | None = None + model_arch: str | None = None + revision: str | None = None + tokenizer_revision: str | None = None + code_revision: str | None = None + seed: int | None = None + logits_processors: list[str | type] | None = None + trust_remote_code: bool = False + dtype: Any = "auto" + attention_backend: Any = None + moe_backend: str = "auto" + hf_overrides: Any = None + limit_mm_per_prompt: dict[str, Any] | None = None active_stream_window: int = Field(default=0, ge=0) duplex_max_sessions: int = Field(default=1, ge=1) enable_sleep_mode: bool = False @@ -267,6 +331,7 @@ class OmniStageModelConfig: task_type: str | None = None codec_frame_rate_hz: float | None = None 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 @@ -281,8 +346,11 @@ class OmniStageModelConfig: @config class OmniStageLoadConfig: - """Per-stage loading behavior.""" + """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 @@ -297,9 +365,11 @@ class OmniStageCacheConfig: vLLM-Omni diffusion-specific cache backends such as TeaCache and Cache-DiT. """ - gpu_memory_utilization: float = Field(default=0.90, gt=0.0, le=1.0) - enable_prefix_caching: bool = False - disable_hybrid_kv_cache_manager: bool = False + kv_cache_memory_bytes: int | None = Field(default=None, ge=0) + # None preserves backend-owned defaults; explicit values still project. + gpu_memory_utilization: float | None = Field(default=None, gt=0.0, le=1.0) + enable_prefix_caching: bool | None = None + disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = Field(default=None, ge=0.0) @@ -307,14 +377,18 @@ class OmniStageCacheConfig: class OmniStageSchedulerConfig: """Per-stage request scheduling behavior.""" - max_num_seqs: int = Field(default=128, ge=1) + 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) - enable_chunked_prefill: bool = False - async_scheduling: bool = True + enable_chunked_prefill: bool | None = None + async_scheduling: bool | None = None def __post_init__(self) -> None: - if self.max_num_batched_tokens is not None and self.max_num_batched_tokens < self.max_num_seqs: + if ( + self.max_num_batched_tokens is not None + and self.max_num_seqs is not None + and self.max_num_batched_tokens < self.max_num_seqs + ): raise ValueError( f"max_num_batched_tokens ({self.max_num_batched_tokens}) must be >= max_num_seqs ({self.max_num_seqs})" ) @@ -322,8 +396,10 @@ class OmniStageSchedulerConfig: @config class OmniStageConnectorConfig: - """Per-stage inter-stage connector wiring.""" + """Per-stage connector wiring and resolved transfer mode.""" + async_chunk: bool = False + omni_kv_config: dict[str, Any] | None = None stage_connector: dict[str, Any] = field( default_factory=lambda: { "name": "SharedMemoryConnector", @@ -336,8 +412,10 @@ class OmniStageConnectorConfig: @config class OmniStageRuntimeConfig: - """Per-stage process placement and runtime behavior.""" + """Per-stage process placement and backend runtime behavior.""" + distributed_executor_backend: Any = None + worker_cls: str | None = None devices: str | None = None num_replicas: int = Field(default=1, ge=1) env: dict[str, Any] | None = None @@ -448,7 +526,12 @@ class _DiffusionConfigProjection: stage_id: int = 0 model: str | None = None model_class_name: str | None = None + engine_backend: str | type = "default" + diffusion_model_runner_cls: str | type | None = None + request_batch_max_wait_ms: float = 0.0 + streaming_output: bool = False model_arch: str | None = None + task_type: str | None = None dtype: Any = "auto" trust_remote_code: bool = False revision: str | None = None @@ -456,6 +539,7 @@ class _DiffusionConfigProjection: dist_timeout: int | None = None nccl_port: int | None = None master_port: int | None = None + scheduler_port: int | None = None host: str | None = None port: int | None = None model_config: dict[str, Any] = field(default_factory=dict) @@ -467,6 +551,7 @@ class _DiffusionConfigProjection: enable_cache_dit_summary: bool = False enable_prompt_embed_cache: bool = False prompt_embed_cache_size: int = Field(default=32, ge=1) + enable_session_state_manager: bool = False diffusion_load_format: str = "default" diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) @@ -479,6 +564,7 @@ class _DiffusionConfigProjection: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_resident_layers: int = Field(default=0, ge=0) pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) @@ -501,6 +587,7 @@ class _DiffusionConfigProjection: step_execution: bool = False supports_multimodal_inputs: bool = False max_multimodal_image_inputs: int | None = None + supports_mixed_reference_inputs: bool = False model_paths: dict[str, str] = field(default_factory=dict) model_loaded: dict[str, bool] = field( default_factory=lambda: { @@ -678,6 +765,7 @@ _DIFFUSION_SHARED_CONFIG_FIELDS = frozenset( "stage_id", "model", "model_arch", + "task_type", "dtype", "trust_remote_code", "revision", @@ -693,6 +781,7 @@ _DIFFUSION_RUNTIME_CONFIG_FIELDS = frozenset( "port", "nccl_port", "master_port", + "scheduler_port", "worker_extension_cls", "enable_stage_verification", "prompt_file_path", @@ -740,21 +829,85 @@ _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) +_CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) _PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__) _PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} +_COMMON_STAGE_ENGINE_FIELDS = ( + _QUANTIZATION_ENGINE_FIELDS + | _MODEL_ENGINE_FIELDS + | _LOAD_ENGINE_FIELDS + | _CACHE_ENGINE_FIELDS + | _SCHEDULER_ENGINE_FIELDS + | _CONNECTOR_ENGINE_FIELDS + | _RUNTIME_ENGINE_FIELDS +) +_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 + | _DIFFUSION_STAGE_ENGINE_FIELDS + | {"parallel_config"} +) + +# This all-stage union is used only to discover global CLI candidates. Stage +# validation must use the execution-type-specific sets below. +_STAGE_ENGINE_FIELDS = _LLM_STAGE_ENGINE_FIELDS | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS +_STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE = { + StageExecutionType.LLM_AR: _LLM_STAGE_ENGINE_FIELDS, + StageExecutionType.LLM_GENERATION: _LLM_STAGE_ENGINE_FIELDS, + StageExecutionType.DIFFUSION: _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS, +} +_PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = { + StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS, + StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS, + StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS, +} + + +def _validate_stage_engine_override_ownership( + stage_id: int, + execution_type: StageExecutionType, + overrides: Mapping[str, Any], +) -> None: + try: + owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] + parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_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): + 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( + f"Stage {stage_id} ({execution_type.value}) has explicit engine argument(s) " + f"with no structured config owner: {names}" + ) def _global_stage_cli_fields() -> frozenset[str]: # Lazy import avoids vllm_omni.config -> omni_config -> engine.arg_utils -> # vllm_omni.config during package-level config imports. - from vllm_omni.engine.arg_utils import OmniEngineArgs + from vllm_omni.engine.arg_utils import OmniEngineArgs, orchestrator_field_names - return ( + candidates = ( frozenset(f.name for f in fields(OmniEngineArgs)) | frozenset(_STAGE_DEPLOY_ENGINE_FIELDS) | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS) - ) - {"model", "stage_id", "stage_configs_path", "async_chunk"} + ) + externally_consumed = ( + _NON_STAGE_ENGINE_CLI_FIELDS + | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig)) + | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS) + ) + return candidates - externally_consumed def _mapping_or_empty(value: Any) -> Mapping[str, Any]: @@ -774,17 +927,33 @@ def _stage_engine_overrides(stage_deploy: StageDeployConfig | None) -> dict[str, value = getattr(stage_deploy, name) if value is not None: overrides[name] = _copy_value(value) - overrides.update(_copy_value(stage_deploy.engine_extras)) + overrides.update( + { + name: _copy_value(value) + for name, value in stage_deploy.engine_extras.items() + if name not in _LEGACY_STAGE_METADATA_EXTRA_FIELDS + } + ) return overrides def _stage_engine_values( stage_deploy: StageDeployConfig | None, + topology: StagePipelineConfig, stage_cli_overrides: Mapping[str, Any] | None = None, ) -> _StageEngineValues: engine = _stage_engine_overrides(stage_deploy) + # Preserve legacy ordering: topology-owned KV roles override deploy + # extras, while an explicit CLI override remains highest priority. + if topology.omni_kv_config: + engine["omni_kv_config"] = _copy_value(topology.omni_kv_config) if stage_cli_overrides: engine.update(_copy_value(stage_cli_overrides)) + _validate_stage_engine_override_ownership( + topology.stage_id, + topology.execution_type, + engine, + ) return _StageEngineValues( quantization=cast( _QuantizationEngineOverrides, @@ -797,6 +966,10 @@ def _stage_engine_values( _SchedulerEngineOverrides, _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS), ), + connector=cast( + _ConnectorEngineOverrides, + _select_engine_overrides(engine, _CONNECTOR_ENGINE_FIELDS), + ), 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), @@ -892,9 +1065,10 @@ class BaseVllmOmniStageConfig: @property def scheduler_cls(self) -> str | None: - return _resolve_scheduler_path( + async_scheduling = self.scheduler_config.async_scheduling + return self.stage_pipeline_config.scheduler_cls or _resolve_scheduler_path( self.stage_pipeline_config.execution_type, - self.scheduler_config.async_scheduling, + True if async_scheduling is None else async_scheduling, ) @property @@ -956,11 +1130,14 @@ StageConfigType: TypeAlias = VllmOmniARStageConfig | VllmOmniGenerationStageConf def _build_common_stage_config_kwargs( + pipeline: PipelineConfig, deploy: DeployConfig, topology: StagePipelineConfig, stage_deploy: StageDeployConfig | None, engine: _StageEngineValues, parallel_config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig, + *, + model: str | None, ) -> tuple[dict[str, Any], str | None, str | None]: input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk)) quantization_config = _build_quantization_config(deploy, engine.quantization) @@ -970,16 +1147,36 @@ def _build_common_stage_config_kwargs( { "stage_pipeline_config": topology, "model_config": _build_model_config( + pipeline, + deploy, topology, stage_deploy, engine.model, duplex_max_sessions=(deploy.duplex_session.max_sessions if deploy.session_mode == "duplex" else 1), + model=model, + ), + "load_config": _build_load_config(topology, engine.load), + "cache_config": _build_cache_config( + deploy, + engine.cache, + topology.execution_type, + ), + "scheduler_config": _build_scheduler_config( + deploy, + engine.scheduler, + topology.execution_type, + ), + "connector_config": _build_connector_config( + deploy, + stage_deploy, + engine.connector, + ), + "runtime_config": _build_runtime_config( + deploy, + stage_deploy, + engine.runtime, + parallel_config, ), - "load_config": _build_load_config(engine.load), - "cache_config": _build_cache_config(deploy, engine.cache), - "scheduler_config": _build_scheduler_config(deploy, engine.scheduler), - "connector_config": _build_connector_config(stage_deploy), - "runtime_config": _build_runtime_config(stage_deploy, engine.runtime, parallel_config), "parallel_config": parallel_config, "quantization_config": _copy_value(quantization_config), }, @@ -1008,10 +1205,12 @@ def _build_ar_stage_config( model: str | None, ) -> VllmOmniARStageConfig: common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + pipeline, deploy, topology, stage_deploy, engine, + model=model, ) return cast( VllmOmniARStageConfig, @@ -1033,10 +1232,12 @@ def _build_generation_stage_config( model: str | None, ) -> VllmOmniGenerationStageConfig: common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + pipeline, deploy, topology, stage_deploy, engine, + model=model, ) return cast( VllmOmniGenerationStageConfig, @@ -1058,18 +1259,20 @@ def _build_diffusion_stage_config( model: str | None, ) -> VllmOmniDiffusionStageConfig: common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + pipeline, deploy, topology, stage_deploy, engine, OmniStageDiffusionParallelConfig, + model=model, ) common_kwargs["diffusion_config"] = _build_diffusion_config_projection( pipeline, deploy, topology, engine.diffusion, - model=model, + model=common_kwargs["model_config"].model, quantization_config=common_kwargs["quantization_config"], ) return cast( @@ -1124,14 +1327,28 @@ def _build_quantization_config( def _build_model_config( + pipeline: PipelineConfig, + deploy: DeployConfig, topology: StagePipelineConfig, stage_deploy: StageDeployConfig | None, engine: _ModelEngineOverrides, *, duplex_max_sessions: int, + model: str | None, ) -> OmniStageModelConfig: default_sampling_params = _stage_sampling_params(stage_deploy, topology) kwargs = _config_kwargs(engine) + kwargs["model"] = _first_defined(kwargs.get("model"), model) + if "model_arch" not in kwargs: + kwargs["model_arch"] = topology.model_arch or pipeline.model_arch or None + if "trust_remote_code" not in kwargs and deploy.trust_remote_code is not None: + kwargs["trust_remote_code"] = _copy_value(deploy.trust_remote_code) + if "dtype" not in kwargs and deploy.dtype is not None: + kwargs["dtype"] = _copy_value(deploy.dtype) + if "active_stream_window" not in kwargs: + kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window) + if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None: + kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir) if "has_sampling_extra_args" not in kwargs: kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args")) if "model_subdir" not in kwargs and topology.model_subdir is not None: @@ -1145,45 +1362,67 @@ def _build_model_config( ) -def _build_load_config(engine: _LoadEngineOverrides) -> OmniStageLoadConfig: - return OmniStageLoadConfig(**_config_kwargs(engine)) +def _build_load_config( + topology: StagePipelineConfig, + engine: _LoadEngineOverrides, +) -> OmniStageLoadConfig: + kwargs = _config_kwargs(engine) + if "skip_mm_profiling" not in kwargs and not topology.requires_multimodal_data: + kwargs["skip_mm_profiling"] = True + return OmniStageLoadConfig(**kwargs) def _build_cache_config( deploy: DeployConfig, engine: _CacheEngineOverrides, + execution_type: StageExecutionType, ) -> OmniStageCacheConfig: 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: + # Match the generation-stage override applied by legacy finalization. + kwargs["disable_hybrid_kv_cache_manager"] = True return OmniStageCacheConfig(**kwargs) def _build_scheduler_config( deploy: DeployConfig, engine: _SchedulerEngineOverrides, + execution_type: StageExecutionType, ) -> OmniStageSchedulerConfig: 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: + kwargs["async_scheduling"] = True return OmniStageSchedulerConfig(**kwargs) -def _build_connector_config(stage_deploy: StageDeployConfig | None) -> OmniStageConnectorConfig: +def _build_connector_config( + deploy: DeployConfig, + stage_deploy: StageDeployConfig | None, + engine: _ConnectorEngineOverrides, +) -> OmniStageConnectorConfig: output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None return OmniStageConnectorConfig( + async_chunk=bool(deploy.async_chunk), + omni_kv_config=_copy_value(engine.get("omni_kv_config")), output_connectors=_copy_value(output_connectors) if output_connectors else None, input_connectors=_copy_value(input_connectors) if input_connectors else None, ) def _build_runtime_config( + deploy: DeployConfig, stage_deploy: StageDeployConfig | None, engine: _RuntimeEngineOverrides, parallel_config: OmniStageParallelConfig, ) -> OmniStageRuntimeConfig: kwargs = _config_kwargs(engine) + if "distributed_executor_backend" not in kwargs and deploy.distributed_executor_backend is not None: + kwargs["distributed_executor_backend"] = _copy_value(deploy.distributed_executor_backend) if "devices" not in kwargs and stage_deploy is not None and stage_deploy.devices is not None: kwargs["devices"] = _copy_value(stage_deploy.devices) if "num_replicas" not in kwargs and stage_deploy is not None: @@ -1233,13 +1472,15 @@ def _build_diffusion_config_projection( topology.model_arch, pipeline.model_arch, ) + if "model_class_name" not in diffusion_kwargs and topology.model_arch is not None: + diffusion_kwargs["model_class_name"] = _copy_value(topology.model_arch) if "dtype" not in diffusion_kwargs and deploy.dtype is not None: diffusion_kwargs["dtype"] = _copy_value(deploy.dtype) if "trust_remote_code" not in diffusion_kwargs and deploy.trust_remote_code is not None: diffusion_kwargs["trust_remote_code"] = _copy_value(deploy.trust_remote_code) if "distributed_executor_backend" not in diffusion_kwargs and deploy.distributed_executor_backend is not None: diffusion_kwargs["distributed_executor_backend"] = _copy_value(deploy.distributed_executor_backend) - if model is not None: + if "model" not in diffusion_kwargs and model is not None: diffusion_kwargs["model"] = model if quantization_config is not None: diffusion_kwargs["quantization_config"] = _copy_value(quantization_config) @@ -1301,6 +1542,7 @@ class VllmOmniConfig: deploy_by_id.get(topology.stage_id), _stage_engine_values( deploy_by_id.get(topology.stage_id), + topology, _stage_cli_overrides(topology.stage_id, cli_overrides), ), model=model, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 4831e6f92..602ca1e20 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -413,6 +413,9 @@ class StageDeployConfig: enable_cpu_offload: bool | None = None enable_layerwise_offload: bool | None = None + enable_distributed_layerwise_offload: bool | None = None + dlo_use_allgather: bool | None = None + dlo_resident_layers: int | None = None # Diffusion-specific debug and observability knobs. enable_diffusion_pipeline_profiler: bool | None = None diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 2e30877ca..68e152063 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -15,6 +15,25 @@ class AttentionBackend(ABC): accept_output_buffer: bool = False supports_piecewise_spans: bool = False + # The backend can represent a contiguous valid K/V prefix by slicing the + # tensors instead of materializing a padding mask. Models may use this to + # avoid a slower masked-attention plan when tail padding is not semantic. + supports_prefix_kv_slicing: bool = False + + # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted. + # Platform resolution rejects an explicit selection outside this set, so a + # hardware-specific backend fails before the model is built. + supported_platforms: tuple[str, ...] | None = None + + @classmethod + def validate_available(cls) -> None: + """Raise if this backend's optional dependencies are missing. + + Called during platform resolution, i.e. before model construction, so a + backend that probes its kernel package lazily still reports the problem + while the user can still act on it. + """ + return None @classmethod def supports_attention_mask(cls) -> bool: @@ -59,6 +78,23 @@ class QueryRange: global_start: int +@dataclass(frozen=True, slots=True) +class VideoTokenLayout: + """Where the video segment sits inside a packed multimodal sequence. + + A model that packs its sequence as ``[prefix | t*h*w video rows | padding]`` + publishes this so backends can recover spatiotemporal locality; the prefix + holds everything that is not video (text, visual conditions, audio). + Publishing it also asserts that any ``attn_mask`` masks only the trailing + padding, so ``prefix_len + t*h*w`` is the used length of the sequence. + + Plain ints, so reading it never forces a device-to-host sync. + """ + + prefix_len: int + latent_grid: tuple[int, int, int] + + @dataclass class AttentionMetadata: attn_mask: torch.Tensor | None = None @@ -83,12 +119,18 @@ class AttentionMetadata: # variable-length query/key sequences for FlashAttention. # "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with # the packed cu_seqlens tensors. + # "valid_kv_length": int — contiguous valid K/V prefix length for a + # backend that advertises supports_prefix_kv_slicing. # Piecewise attention metadata (mixed causal/full masks). # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention. full_attn_spans: list[list[tuple[int, int]]] | None = None query_ranges: tuple[QueryRange, ...] | None = None + # Geometry of the video segment for backends that exploit spatiotemporal + # locality (block-sparse selection, tiled masks). Dense backends ignore it. + video_layout: VideoTokenLayout | None = None + T = TypeVar("T", bound=AttentionMetadata) diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index 78e428796..ffcc2fb44 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -17,6 +17,7 @@ logger = init_logger(__name__) class CuDNNAttentionBackend(AttentionBackend): accept_output_buffer: bool = True + supports_prefix_kv_slicing: bool = True @classmethod def supports_attention_mask(cls) -> bool: @@ -61,7 +62,26 @@ class CuDNNAttentionImpl(AttentionImpl): ) -> torch.Tensor: attention_mask = None if attn_metadata: - attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k") + valid_kv_length = attn_metadata.extra.get("valid_kv_length") + if attn_metadata.attn_mask is None and isinstance(valid_kv_length, int): + if not 0 < valid_kv_length <= key.shape[1]: + raise ValueError( + "valid_kv_length must be within the K/V sequence, " + f"got {valid_kv_length} for length {key.shape[1]}" + ) + # A contiguous valid prefix is mathematically equivalent to a + # broadcast key-padding mask. Slicing K/V keeps Q/output in the + # checkpoint's aligned layout while letting cuDNN select its + # much faster mask-free FMHA plan. + key = key[:, :valid_kv_length] + value = value[:, :valid_kv_length] + else: + attention_mask = _maybe_reshape_attn_mask( + query, + key, + attn_metadata.attn_mask, + mask_mode="broadcast_k", + ) enable_gqa = query.shape[2] != key.shape[2] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py new file mode 100644 index 000000000..10ab55991 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.model_executor.models.utils import extract_layer_index + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend +from vllm_omni.diffusion.config import get_current_diffusion_config_or_none +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + +logger = init_logger(__name__) + +# The rf_v2 kernel only implements a 128-token block. +_BLOCK_SIZE = 128 + +# Below this many video blocks, the pooling and gather that block selection adds +# cost more than the QK work it removes, so stay dense. +_MIN_VIDEO_BLOCKS = 32 + +# rf_v2's ``input_layout`` describes the caller's tensors, and everything below +# slices the sequence on dim 1. vLLM-Omni diffusion attention hands the impl +# [B, S, N, D], so that is the only layout this backend accepts. +_INPUT_LAYOUT = "BSND" + +# rf_v2 precision mode: 0 = high precision, 1 = high performance. Kept at the +# precise setting because the sparsity knob is the intended perf lever, and on +# A5 devices the kernel is routed to rf_v3, which overrides this anyway. +_INNER_PRECISE = 0 + +_WRONG_PLATFORM = ( + "RAINFUSION_ATTN runs the MindIE-SD rf_v2 kernel and is available on Ascend NPU only. " + "Select FLASH_ATTN or TORCH_SDPA on this platform." +) + +_MISSING_MINDIESD = ( + "RAINFUSION_ATTN requires MindIE-SD. Please install MindIE-SD to enable RainFusion sparse " + "attention on Ascend NPU. For installation details, see https://gitcode.com/Ascend/MindIE-SD " + "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN" +) + + +def _try_extract_layer_index(prefix: str) -> int | None: + if not prefix: + return None + try: + return extract_layer_index(prefix) + except (AssertionError, ValueError): + return None + + +@dataclass(frozen=True) +class RainFusionConfig: + """Resolved RainFusion controls for one attention layer. + + ``sparsity`` is the nominal fraction of key blocks dropped per query block. + The realized sparsity is lower because rf_v2 always keeps the prefix rows and + the first-frame blocks. ``start_step`` and ``skip_layers`` are the accuracy + knobs: early denoise steps and specific DiT blocks stay dense. + """ + + sparsity: float = 0.0 + start_step: int = 0 + skip_layers: frozenset[int] = frozenset() + + @classmethod + def from_backend_kwargs(cls, backend_kwargs: dict | None) -> RainFusionConfig: + bk = backend_kwargs or {} + return cls( + sparsity=float(bk.get("sparsity", 0.0)), + start_step=int(bk.get("start_step", 0)), + skip_layers=frozenset(bk.get("skip_layers") or ()), + ) + + @property + def enabled(self) -> bool: + return self.sparsity > 0.0 + + +@dataclass(frozen=True) +class RainFusionPlan: + """Per-forward geometry handed to the rf_v2 kernel.""" + + prefix_len: int + used_len: int + latent_shape: list[int] + + +class RainFusionAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + supported_platforms: tuple[str, ...] = ("npu",) + + @classmethod + def validate_available(cls) -> None: + from importlib.util import find_spec + + if find_spec("mindiesd") is None: + raise ValueError(_MISSING_MINDIESD) + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 96, 128, 192, 256] + + @staticmethod + def get_name() -> str: + return "RAINFUSION_ATTN" + + @staticmethod + def get_impl_cls() -> type[RainFusionAttentionImpl]: + return RainFusionAttentionImpl + + +class RainFusionAttentionImpl(AttentionImpl): + """Block-sparse video attention via MindIE-SD RainFusion (rf_v2) on Ascend NPU. + + Sparsity applies only to the video segment of a packed multimodal sequence, + whose extent the model publishes as ``AttentionMetadata.video_layout``. Every + other case — warmup denoise steps, exempt layers, a layer that does not declare + ``qkv_layout="BSND"``, sequences without a published video segment, video + segments too short to pay for block selection or not aligned to the kernel's + block size — delegates to FlashAttention, so a model can select this backend + unconditionally. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + qkv_layout: str | None = None, + backend_kwargs: dict[str, Any] | None = None, + **extra_impl_args, + ) -> None: + self.num_heads = num_heads + self.causal = causal + self.softmax_scale = softmax_scale + self.qkv_layout = qkv_layout + + self.rainfusion = RainFusionConfig.from_backend_kwargs(backend_kwargs) + self.layer_idx = _try_extract_layer_index(prefix) + + if self.rainfusion.enabled: + self._validate_parallel_config() + if causal: + raise ValueError( + "RAINFUSION_ATTN does not support causal attention: rf_v2 selects key " + "blocks by pooled relevance and cannot express a causal mask. Select " + "FLASH_ATTN for causal roles." + ) + if qkv_layout is not None and qkv_layout.upper() != _INPUT_LAYOUT: + raise ValueError( + f"RAINFUSION_ATTN needs {_INPUT_LAYOUT} tensors to locate the video segment along " + f"the sequence axis, but this layer declares qkv_layout={qkv_layout!r}. Select " + "FLASH_ATTN for this role." + ) + + self.dense_fallback = FlashAttentionBackend.get_impl_cls()( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + prefix=prefix, + qkv_layout=qkv_layout, + ) + + def _validate_parallel_config(self) -> None: + config = get_current_diffusion_config_or_none() + parallel_config = getattr(config, "parallel_config", None) + ring_degree = getattr(parallel_config, "ring_degree", 1) + if ring_degree > 1: + # Ring gives each rank a slice of the sequence, so block selection + # would score only local keys and the layer bypasses the backend + # entirely (see Attention._run_ring_attention). + raise ValueError( + "RAINFUSION_ATTN is not compatible with ring sequence parallelism " + f"(ring_degree={ring_degree}): rf_v2 needs the whole key sequence to rank " + "blocks. Use Ulysses SP (ring_degree=1) instead." + ) + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + # ROCm and MUSA route through forward_cuda by default, so this covers them too. + raise NotImplementedError(_WRONG_PLATFORM) + + def forward_xpu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + raise NotImplementedError(_WRONG_PLATFORM) + + def forward_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + plan = self._resolve_plan(attn_metadata) + if plan is None: + return self.dense_fallback.forward_npu(query, key, value, attn_metadata) + return self._forward_sparse_npu(query, key, value, plan) + + def _resolve_plan(self, attn_metadata: AttentionMetadata | None) -> RainFusionPlan | None: + """Return the rf_v2 geometry, or None when this forward must stay dense.""" + rf = self.rainfusion + if not rf.enabled: + return None + if self.layer_idx is not None and self.layer_idx in rf.skip_layers: + return None + if is_forward_context_available(): + step_idx = get_forward_context().denoise_step_idx + if step_idx is not None and step_idx < rf.start_step: + return None + if self.qkv_layout is None: + # The sparse path reads the sequence off dim 1, which the tensors alone + # do not establish, and the dense fallback resolves an absent layout its + # own way. Sparsifying on an assumption would put the two paths on + # different axes, so an undeclared layout stays dense. + logger.warning_once( + "RAINFUSION_ATTN staying dense: this layer does not declare qkv_layout, and rf_v2 " + "needs %s to locate the video segment along the sequence axis. Set qkv_layout=%r on " + "the Attention layer to enable sparsity.", + _INPUT_LAYOUT, + _INPUT_LAYOUT, + ) + return None + + if attn_metadata is None: + return None + + layout = attn_metadata.video_layout + if layout is None: + logger.warning_once( + "RAINFUSION_ATTN staying dense: this attention role carries no video segment. The " + "model must publish AttentionMetadata.video_layout for the sequence to be sparsified." + ) + return None + max_seqlen_q = attn_metadata.extra.get("max_seqlen_q") + if max_seqlen_q is None: + logger.warning_once( + "RAINFUSION_ATTN staying dense: attention metadata is missing max_seqlen_q, so the " + "video segment cannot be confirmed to be the tail of packed document 0." + ) + return None + + prefix_len = int(layout.prefix_len) + latent_shape = [int(dim) for dim in layout.latent_grid] + # rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of + # the packed sequence holds those rows; anything past it is alignment + # padding that rf_v2 must not see. + video_len = math.prod(latent_shape) + used_len = prefix_len + video_len + + if used_len != int(max_seqlen_q): + logger.warning_once( + "RAINFUSION_ATTN staying dense: prefix (%d) plus latent grid %s does not fill " + "packed document 0 (%d rows). rf_v2 requires the video segment to be its tail.", + prefix_len, + tuple(latent_shape), + int(max_seqlen_q), + ) + return None + if video_len < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE: + logger.warning_once( + "RAINFUSION_ATTN staying dense: %d video rows is under the %d-row " + "(%d block) threshold where sparse selection pays off.", + video_len, + _MIN_VIDEO_BLOCKS * _BLOCK_SIZE, + _MIN_VIDEO_BLOCKS, + ) + return None + if video_len % _BLOCK_SIZE != 0: + # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but + # builds its block mask from the two segments pooled apart. The counts + # only agree when the video segment ends on a block boundary; otherwise + # a prefix block is forced on out of range, and the block straddling the + # seam holds real prefix rows that selection may drop. Padding the + # sequence to realign it is not an option because rf_v2 ignores + # attn_mask, so pad keys would take a share of every softmax denominator. + logger.warning_once( + "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a " + "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a " + "multiple of %d to enable sparsity.", + tuple(latent_shape), + video_len, + _BLOCK_SIZE, + _BLOCK_SIZE, + ) + return None + + logger.info_once( + "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, " + "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal " + "because prefix and first-frame blocks are always kept.", + rf.sparsity, + rf.start_step, + len(rf.skip_layers), + tuple(latent_shape), + prefix_len, + video_len, + ) + return RainFusionPlan( + prefix_len=prefix_len, + used_len=used_len, + latent_shape=latent_shape, + ) + + def _forward_sparse_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + plan: RainFusionPlan, + ) -> torch.Tensor: + try: + from mindiesd import sparse_attention + except ImportError: + raise ImportError(_MISSING_MINDIESD) + + used = plan.used_len + q, k, v = (tensor[:, :used] for tensor in (query, key, value)) + # Ulysses has already gathered the full sequence onto this rank and split + # the heads, so read the head count off the tensor rather than num_heads. + out = sparse_attention( + q, + k, + v, + scale=self.softmax_scale, + head_num=query.shape[-2], + input_layout=_INPUT_LAYOUT, + inner_precise=_INNER_PRECISE, + sparse_type="rf_v2", + txt_len=plan.prefix_len, + block_size=_BLOCK_SIZE, + latent_shape_q=plan.latent_shape, + latent_shape_k=plan.latent_shape, + sparsity=self.rainfusion.sparsity, + ) + if used == query.shape[1]: + return out + padded = torch.zeros_like(query) + padded[:, :used] = out + return padded diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index 92d3b6926..fdfaf5f9c 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -65,6 +65,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta): FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend" FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend" TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend" + RAINFUSION_ATTN = "vllm_omni.diffusion.attention.backends.rainfusion_attn.RainFusionAttentionBackend" def get_path(self, include_classname: bool = True) -> str: """Get the class path for this backend (respects overrides). diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py index 718484a59..c63280fc8 100644 --- a/vllm_omni/diffusion/attention/backends/trtllm_attn.py +++ b/vllm_omni/diffusion/attention/backends/trtllm_attn.py @@ -180,6 +180,7 @@ class TrtllmAttentionImpl(AttentionImpl): prefix: str = "", qkv_layout: str | None = None, backend_kwargs: dict | None = None, + role: str = "self", **extra_impl_args, ) -> None: self.num_heads = num_heads @@ -187,6 +188,7 @@ class TrtllmAttentionImpl(AttentionImpl): self.softmax_scale = softmax_scale self.causal = causal self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.role = role self.skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs) self._warned_missing_timestep = False @@ -255,8 +257,16 @@ class TrtllmAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: + extra = getattr(attn_metadata, "extra", {}) if attn_metadata is not None else {} + packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k") + present_packed_keys = [key for key in packed_keys if key in extra] + if present_packed_keys and len(present_packed_keys) != len(packed_keys): + missing = sorted(set(packed_keys) - set(present_packed_keys)) + raise ValueError(f"Incomplete packed TRTLLM attention metadata; missing {missing}") + has_packed_metadata = len(present_packed_keys) == len(packed_keys) + attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None - if attn_mask is not None: + if attn_mask is not None and not has_packed_metadata: raise ValueError( "TRTLLM_ATTN does not support attn_mask (Wan sets one only under SP with " "mask_sp_padding=True and a non-divisible seq len). Either set " @@ -270,30 +280,96 @@ class TrtllmAttentionImpl(AttentionImpl): "another backend via --diffusion-attention-backend." ) - batch, q_len, num_q_heads, head_dim = query.shape + physical_batch, q_len, num_q_heads, head_dim = query.shape kv_len, num_kv_heads = key.shape[1], key.shape[2] device = query.device - q = query.reshape(batch * q_len, num_q_heads, head_dim).contiguous() - k = key.reshape(batch * kv_len, num_kv_heads, head_dim).contiguous() - v = value.reshape(batch * kv_len, num_kv_heads, head_dim).contiguous() + q = query.reshape(physical_batch * q_len, num_q_heads, head_dim).contiguous() + k = key.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous() + v = value.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous() + output_tokens = q.shape[0] - seq_lens = torch.full((batch,), kv_len, dtype=torch.int32, device=device) - cu_seq_lens_q = torch.arange(0, (batch + 1) * q_len, step=q_len, dtype=torch.int32, device=device) - cu_seq_lens_kv = torch.arange(0, (batch + 1) * kv_len, step=kv_len, dtype=torch.int32, device=device) + if has_packed_metadata: + cu_seq_lens_q = extra["cu_seqlens_q"] + cu_seq_lens_kv = extra["cu_seqlens_k"] + if cu_seq_lens_q.ndim != 1 or cu_seq_lens_kv.ndim != 1: + raise ValueError("Packed TRTLLM attention cu_seqlens tensors must be one-dimensional") + if cu_seq_lens_q.numel() != cu_seq_lens_kv.numel() or cu_seq_lens_q.numel() < 2: + raise ValueError("Packed TRTLLM attention requires matching non-empty Q and KV sequence batches") + if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]: + raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens") + + # Drop trailing empty sequences retained by packed producers when + # alignment adds no padding tokens. + while ( + cu_seq_lens_q.numel() > 2 + and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item()) + and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item()) + ): + cu_seq_lens_q = cu_seq_lens_q[:-1] + cu_seq_lens_kv = cu_seq_lens_kv[:-1] + + # A packed prefix mask denotes structural suffix padding, not another + # document. Exclude it before quantization and attention; in particular, + # SAGE block quantization must not span the valid/padding boundary. + if attn_mask is not None: + if q.shape[0] != k.shape[0] or attn_mask.numel() != q.shape[0]: + raise ValueError("Packed TRTLLM prefix masks require equal Q/K lengths") + flat_mask = attn_mask.reshape(-1).to(dtype=torch.bool) + valid_tokens = int(flat_mask.sum().item()) + expected_mask = torch.arange(flat_mask.numel(), device=device) < valid_tokens + if not torch.equal(flat_mask, expected_mask): + raise ValueError("TRTLLM_ATTN only supports prefix-valid packed attention masks") + + q_boundary = (cu_seq_lens_q == valid_tokens).nonzero(as_tuple=False) + kv_boundary = (cu_seq_lens_kv == valid_tokens).nonzero(as_tuple=False) + if q_boundary.numel() != 1 or kv_boundary.numel() != 1: + raise ValueError("Packed TRTLLM prefix length must be a Q/K sequence boundary") + q_end = int(q_boundary.item()) + 1 + kv_end = int(kv_boundary.item()) + 1 + if q_end != kv_end: + raise ValueError("Packed TRTLLM prefix masks require matching Q/K sequence batches") + + q = q[:valid_tokens] + k = k[:valid_tokens] + v = v[:valid_tokens] + cu_seq_lens_q = cu_seq_lens_q[:q_end].contiguous() + cu_seq_lens_kv = cu_seq_lens_kv[:kv_end].contiguous() + + batch = cu_seq_lens_q.numel() - 1 + seq_lens = (cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous() + max_q_len = int(extra["max_seqlen_q"]) + max_kv_len = int(extra["max_seqlen_k"]) + else: + batch = physical_batch + seq_lens = torch.full((batch,), kv_len, dtype=torch.int32, device=device) + cu_seq_lens_q = torch.arange(0, (batch + 1) * q_len, step=q_len, dtype=torch.int32, device=device) + cu_seq_lens_kv = torch.arange(0, (batch + 1) * kv_len, step=kv_len, dtype=torch.int32, device=device) + max_q_len = q_len + max_kv_len = kv_len workspace = self._get_workspace(device) bmm1_scale = self.softmax_scale bmm2_scale = 1.0 - _skip_factor = self._resolve_skip_factor(kv_len) + _skip_factor = self._resolve_skip_factor(max_kv_len) # SAGE kwargs are only understood by newer FlashInfer builds; pass them exclusively when # SAGE quant is active (which already requires the kernel, checked at init) so the dense # path stays compatible with older builds that lack these parameters. sage_kwargs: dict = {} - if self.quant.enabled: + # The SAGE kernel requires every KV sequence to contain at least one full + # quantization block. Small auxiliary attention sites use the dense kernel. + use_sage = self.quant.enabled and bool(torch.all(seq_lens >= self.quant.k_block_size).item()) + if self.quant.enabled and not use_sage: + message = ( + f"TRTLLM_ATTN SAGE quantization is configured for attention role {self.role!r}, but at least one " + f"KV sequence is shorter than k_block_size={self.quant.k_block_size}. Falling back to dense " + "attention for this input." + ) + logger.warning_once(message) + if use_sage: q, k, v, sage_attn_sfs, sage_block_sizes = self.quant.quantize(q, k, v, self._sage_quantize_fn) sage_kwargs["sage_attn_sfs"] = sage_attn_sfs sage_kwargs["num_elts_per_sage_attn_blk"] = sage_block_sizes @@ -304,8 +380,8 @@ class TrtllmAttentionImpl(AttentionImpl): value=v, workspace_buffer=workspace, seq_lens=seq_lens, - max_q_len=q_len, - max_kv_len=kv_len, + max_q_len=max_q_len, + max_kv_len=max_kv_len, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, o_sf_scale=-1.0, @@ -319,4 +395,12 @@ class TrtllmAttentionImpl(AttentionImpl): skip_softmax_threshold_scale_factor=_skip_factor, **sage_kwargs, ) - return out.reshape(batch, q_len, num_q_heads, head_dim) + if out.shape[0] != output_tokens: + padded_out = torch.zeros( + (output_tokens, num_q_heads, head_dim), + dtype=out.dtype, + device=out.device, + ) + padded_out[: out.shape[0]] = out + out = padded_out + return out.reshape(physical_batch, q_len, num_q_heads, head_dim) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 77d65f79e..7eb74a941 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -89,6 +89,15 @@ class Attention(nn.Module): role_category=role_category, allow_trtllm_default=allow_trtllm_default, ) + parallel_config = getattr(config, "parallel_config", None) + allgather_degree = getattr(parallel_config, "allgather_degree", 1) + # TODO: Move AllGather-KV compatibility into an AttentionBackend capability + # so validation does not depend on backend names. + if not skip_sequence_parallel and allgather_degree > 1 and attn_backend_cls.get_name() == "TRTLLM_ATTN": + raise ValueError( + "TRTLLM_ATTN does not support AllGather-KV sequence parallelism. " + "Set --allgather-degree 1 or select another diffusion attention backend." + ) if spec is not None: backend_kwargs = spec.backend_kwargs() self.backend_pref = spec.backend @@ -107,6 +116,7 @@ class Attention(nn.Module): qkv_layout=qkv_layout, prefix=prefix, backend_kwargs=backend_kwargs, + role=role, ) # Instantiate fallback backend for float32 support self.sdpa_fallback = SDPABackend.get_impl_cls()( diff --git a/vllm_omni/diffusion/cache/cachedit/backend.py b/vllm_omni/diffusion/cache/cachedit/backend.py index ac6f93785..c30a888a8 100644 --- a/vllm_omni/diffusion/cache/cachedit/backend.py +++ b/vllm_omni/diffusion/cache/cachedit/backend.py @@ -3,10 +3,13 @@ """Generic Cache-DiT backend lifecycle and integration.""" +from __future__ import annotations + from collections.abc import Callable -from typing import Any, TypeAlias +from typing import TYPE_CHECKING, Any, TypeAlias, cast import cache_dit +import torch.nn as nn from cache_dit import BlockAdapter, DBCacheConfig from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter from vllm.logger import init_logger @@ -18,6 +21,9 @@ from vllm_omni.diffusion.cache.cachedit.config import ( ) from vllm_omni.diffusion.data import DiffusionCacheConfig +if TYPE_CHECKING: + from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + logger = init_logger(__name__) RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None] @@ -27,14 +33,18 @@ CacheDiTEnabler: TypeAlias = Callable[[Any, DiffusionCacheConfig], RefreshCacheC CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {} -def cache_summary(pipeline: Any, details: bool = True) -> None: +def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]: + """Return the pipeline DiT attributes that are present at runtime.""" + names = getattr(pipeline, "_dit_modules", None) + if not isinstance(names, (list, tuple)): + names = ("transformer",) + return tuple(name for name in names if isinstance(name, str) and getattr(pipeline, name, None) is not None) + + +def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None: """Log Cache-DiT statistics for every transformer on the pipeline.""" - transformers = [ - transformer - for attribute in ("transformer", "transformer_2") - if (transformer := getattr(pipeline, attribute, None)) is not None - ] + transformers = [getattr(pipeline, name) for name in _dit_module_names(pipeline)] for transformer in transformers: cache_dit.summary(transformer, details=details) @@ -42,19 +52,34 @@ def cache_summary(pipeline: Any, details: bool = True) -> None: logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute") -def _default_get_pipeline_transformer(pipeline: Any) -> Any: - return pipeline.transformer +def _default_get_pipeline_transformer( + pipeline: SupportsComponentDiscovery, +) -> nn.Module: + return cast(nn.Module, getattr(pipeline, "transformer")) + + +def _make_pipeline_transformer_getter( + name: str, +) -> Callable[[SupportsComponentDiscovery], nn.Module]: + def get_pipeline_transformer( + pipeline: SupportsComponentDiscovery, + ) -> nn.Module: + return cast(nn.Module, getattr(pipeline, name)) + + return get_pipeline_transformer def _build_cache_context_refresh( cache_config: DiffusionCacheConfig, - get_pipeline_transformer: Callable[[Any], Any] = _default_get_pipeline_transformer, + get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, ) -> RefreshCacheContextFunc: """Build the cache context refresh callback for one transformer.""" projected_config = CacheDiTConfig.from_diffusion_config(cache_config) - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + def refresh_cache_context( + pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True + ) -> None: transformer = get_pipeline_transformer(pipeline) # Cache-DiT has no predefined SCM mask for these small step counts. @@ -80,10 +105,11 @@ def _build_cache_context_refresh( def enable_cache_for_dit( - pipeline: Any, + pipeline: SupportsComponentDiscovery, cache_config: DiffusionCacheConfig, block_adapter: BlockAdapter | None = None, adapter_cls: type[CachedAdapter] | None = None, + get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, ) -> RefreshCacheContextFunc: """Enable Cache-DiT for a standard single-transformer DiT pipeline.""" @@ -98,7 +124,7 @@ def enable_cache_for_dit( db_cache_config.max_warmup_steps, ) - transformer = _default_get_pipeline_transformer(pipeline) + transformer = get_pipeline_transformer(pipeline) cache_target = transformer if block_adapter is None else block_adapter if adapter_cls is not None: adapter_cls.apply( @@ -128,13 +154,16 @@ def enable_cache_for_dit( calibrator_config=calibrator_config, ) - return _build_cache_context_refresh(cache_config) + return _build_cache_context_refresh(cache_config, get_pipeline_transformer) -def _maybe_build_block_adapter(pipeline: Any) -> BlockAdapter | None: +def _maybe_build_block_adapter( + pipeline: SupportsComponentDiscovery, + get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, +) -> BlockAdapter | None: """Build the model-declared block adapter, when one is configured.""" - transformer = _default_get_pipeline_transformer(pipeline) + transformer = get_pipeline_transformer(pipeline) adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) if adapter_config is None: logger.info( @@ -160,10 +189,13 @@ def _maybe_build_block_adapter(pipeline: Any) -> BlockAdapter | None: ) -def _maybe_get_cached_adapter_cls(pipeline: Any) -> type[CachedAdapter] | None: +def _maybe_get_cached_adapter_cls( + pipeline: SupportsComponentDiscovery, + get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, +) -> type[CachedAdapter] | None: """Return the custom cached adapter declared by the transformer.""" - transformer = _default_get_pipeline_transformer(pipeline) + transformer = get_pipeline_transformer(pipeline) adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) return None if adapter_config is None else adapter_config.cached_adapter_cls @@ -180,29 +212,37 @@ class CacheDiTBackend(CacheBackend): config = cache_config super().__init__(config) - self._refresh_func: RefreshCacheContextFunc | None = None + self._refresh_funcs: list[RefreshCacheContextFunc] = [] - def enable(self, pipeline: Any) -> None: + def enable(self, pipeline: SupportsComponentDiscovery) -> None: pipeline_name = type(pipeline).__name__ custom_enabler = CUSTOM_DIT_ENABLERS.get(pipeline_name) if custom_enabler is not None: logger.info("Using custom cache-dit enabler for model: %s", pipeline_name) - self._refresh_func = custom_enabler(pipeline, self.config) + self._refresh_funcs = [custom_enabler(pipeline, self.config)] else: - block_adapter = _maybe_build_block_adapter(pipeline) - adapter_cls = _maybe_get_cached_adapter_cls(pipeline) - self._refresh_func = enable_cache_for_dit( - pipeline, - self.config, - block_adapter, - adapter_cls, - ) + self._refresh_funcs = [] + for name in _dit_module_names(pipeline): + get_transformer = _make_pipeline_transformer_getter(name) + block_adapter = _maybe_build_block_adapter(pipeline, get_transformer) + adapter_cls = _maybe_get_cached_adapter_cls(pipeline, get_transformer) + self._refresh_funcs.append( + enable_cache_for_dit( + pipeline, + self.config, + block_adapter, + adapter_cls, + get_transformer, + ) + ) + if not self._refresh_funcs: + raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT") self.enabled = True logger.info("Cache-dit enabled successfully on %s", pipeline_name) - def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - if not self.enabled or self._refresh_func is None: + def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None: + if not self.enabled or not self._refresh_funcs: logger.warning("Cache-dit is not enabled. Cannot refresh cache context.") return @@ -211,7 +251,8 @@ class CacheDiTBackend(CacheBackend): "Refreshing cache context for transformer with num_inference_steps: %s", num_inference_steps, ) - self._refresh_func(pipeline, num_inference_steps, verbose) + for refresh_func in self._refresh_funcs: + refresh_func(pipeline, num_inference_steps, verbose) __all__ = [ diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 24260a5b7..b0ef2a74a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -541,16 +541,15 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de """ from vllm.transformers_utils.config import get_hf_file_to_dict + from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index + if not model: return None is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"} - # Diffusers models: read _class_name from model_index.json. - try: - model_index = get_hf_file_to_dict("model_index.json", model) - except Exception: - model_index = None + # Diffusers models: read _class_name from the pipeline index. + model_index = get_diffusion_model_index(model) if model_index is not None: return model_index.get("_class_name") if diffusion_load_format == "diffusers": @@ -603,6 +602,10 @@ class OmniDiffusionConfig: model_class_name: str | None = None + # Optional model-defined startup task. Pipelines may use this to select + # task-specific components or weights before serving requests. + task_type: str | None = None + dtype: torch.dtype = torch.bfloat16 model_config: dict[str, Any] = field(default_factory=dict) @@ -687,9 +690,13 @@ class OmniDiffusionConfig: # Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1) enable_distributed_layerwise_offload: bool = False # If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires - # concurrent requests in DP mode). If False: each rank loads full weights - # via H2D only (N× CPU memory, but no AllGather synchronization needed). + # concurrent requests in DP mode). If False: each rank streams the standard + # loader's rank-local tensors (including TP-local shards) via H2D only. + # This avoids AllGather synchronization, while host memory follows the + # loader's existing rank-local layout instead of adding a second DP shard. dlo_use_allgather: bool = True + # Leading main-DiT blocks kept resident by distributed layerwise offload. + dlo_resident_layers: int = 0 pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -1139,18 +1146,23 @@ class OmniDiffusionConfig: def enrich_config(self) -> None: """Load model metadata from HuggingFace and populate config fields. - Diffusers-style models expose ``model_index.json`` with ``_class_name``. + Diffusers-style models expose a pipeline index with ``_class_name``. Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``, so we fall back to reading that and mapping model_type manually. """ from vllm.transformers_utils.config import get_hf_file_to_dict + from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index + # Default model_class_name for diffusers adapter if self.model_class_name is None and self.diffusion_load_format == "diffusers": self.model_class_name = "DiffusersAdapterPipeline" try: - config_dict = get_hf_file_to_dict("model_index.json", self.model) + config_dict = get_diffusion_model_index( + self.model, + revision=self.revision, + ) if config_dict is not None: if self.model_class_name is None: self.model_class_name = config_dict.get("_class_name", None) @@ -1165,7 +1177,7 @@ class OmniDiffusionConfig: self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name) except (KeyError, AttributeError) as exc: logger.warning( - "Could not find valid _class_name for diffusers pipeline in model_index.json: %s. " + "Could not find a valid _class_name in the Diffusers pipeline index: %s. " "Without the underlying pipeline class the dummy run may omit required inputs.", exc, ) @@ -1178,16 +1190,16 @@ class OmniDiffusionConfig: else: self.set_tf_model_config(TransformerConfig()) else: - raise FileNotFoundError("model_index.json not found") + raise FileNotFoundError("Diffusers pipeline index not found") except (AttributeError, OSError, ValueError, FileNotFoundError): # Skip transformer config loading for diffusers adapter # (non-DiT models don't have a separate transformer folder/config) if self.diffusion_load_format == "diffusers": self.set_tf_model_config(TransformerConfig()) logger.warning( - "Could not find valid model_index.json per diffusers format. " + "Could not find a valid pipeline index per Diffusers format. " "This model is likely unsupported by the diffusers backend. " - "Also, without knowing the underlying diffusers pipeline class from model_index.json, " + "Also, without knowing the underlying pipeline class from its index, " "the dummy run will input only text prompt, which may cause errors for pipelines " "that require additional inputs." ) @@ -1203,7 +1215,7 @@ class OmniDiffusionConfig: self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() return - raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") + raise ValueError(f"Could not find config.json or a Diffusers pipeline index for {self.model}") self.set_tf_model_config(TransformerConfig.from_dict(cfg)) model_type = cfg.get("model_type") @@ -1430,6 +1442,10 @@ class SkipSoftmaxSpec: if self.target_sparsity is not None and self.threshold is not None: raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.") + @property + def enabled(self) -> bool: + return self.threshold is not None or self.target_sparsity is not None + @dataclass class AttnQuantSpec: @@ -1457,6 +1473,33 @@ class AttnQuantSpec: return self.dtype_qk is not None or self.dtype_vo is not None +# Backends that select key blocks instead of attending densely, and so accept +# a ``block_sparse`` spec. Each maps the same knobs onto its own kernel. +BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"}) + + +@dataclass +class BlockSparseSpec: + """User-facing controls shared by block-sparse attention backends. + + ``sparsity`` is the nominal fraction of key blocks dropped per query block; + ``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an + index selector such as "0-3,38") exempts individual DiT blocks. Those two are + the accuracy knobs to trade back quality at a fixed ``sparsity``. + """ + + sparsity: float = 0.8 + start_step: int = 0 + skip_layers: str | list[int] | None = None + skip_layer_indices: set[int] | None = field(default=None, repr=False) + + def __post_init__(self) -> None: + self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0 + if self.start_step < 0: + raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.") + self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers) + + @dataclass class AttentionSpec: """Specifies a backend and its typed backend-specific config for one attention role.""" @@ -1464,6 +1507,7 @@ class AttentionSpec: backend: str skip_softmax: SkipSoftmaxSpec | None = None quant: AttnQuantSpec | None = None + block_sparse: BlockSparseSpec | None = None skip_calibration: dict | None = field(default=None, repr=False) def __post_init__(self) -> None: @@ -1471,6 +1515,7 @@ class AttentionSpec: raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}") self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax") self.quant = self._coerce(self.quant, AttnQuantSpec, "quant") + self.block_sparse = self._coerce(self.block_sparse, BlockSparseSpec, "block_sparse") if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN": raise ValueError( f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. " @@ -1481,6 +1526,15 @@ class AttentionSpec: f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but " f"backend={self.backend!r}. Remove quant or set a supported backend." ) + if self.backend.upper() in BLOCK_SPARSE_BACKENDS: + # Selecting the backend is the opt-in; without an explicit block the + # defaults apply rather than silently running dense. + self.block_sparse = self.block_sparse or BlockSparseSpec() + elif self.block_sparse is not None: + raise ValueError( + f"block_sparse is only supported by the {sorted(BLOCK_SPARSE_BACKENDS)} backends, but " + f"backend={self.backend!r}. Remove block_sparse or set a supported backend." + ) @staticmethod def _coerce(value: Any, cls: type, field_name: str) -> Any: @@ -1513,6 +1567,12 @@ class AttentionSpec: if q.flashinfer_backend is not None: quant_kw["flashinfer_backend"] = q.flashinfer_backend kw["quant"] = quant_kw + if self.block_sparse is not None: + bs = self.block_sparse + kw["sparsity"] = bs.sparsity + kw["start_step"] = bs.start_step + if bs.skip_layer_indices: + kw["skip_layers"] = sorted(bs.skip_layer_indices) return kw or None @@ -1582,7 +1642,7 @@ class AttentionConfig: normalized[role] = node return - spec_keys = {"backend", "skip_softmax", "quant"} + spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"} node_dict = dict(node) node_keys = set(node_dict) if node_keys & spec_keys: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index e89b95571..086206fa0 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -237,7 +237,7 @@ class DiffusionEngine: if getattr(self.od_config, "parallel_config", None) is not None: dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1) dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) - if dp_size > 1 and dist_offload: + if dp_size > 1 and dist_offload and getattr(self.od_config, "dlo_use_allgather", True): self.scheduler.max_num_running_reqs = dp_size self.dp_concurrent = True # Set batch admission wait so concurrent requests accumulate diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index fa0213fca..fa7130e40 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -416,11 +416,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): ) extra_args_signatures: set = set() for nr in new_reqs: - ea = getattr(nr.req, "extra_args", None) - if ea and isinstance(ea, dict): - extra_args_signatures.add(json.dumps(ea, sort_keys=True)) - else: - extra_args_signatures.add(None) + ea = getattr(nr.req.sampling_params, "extra_args", None) + extra_args_signatures.add(json.dumps(ea, sort_keys=True, default=repr) if ea is not None else None) if len(extra_args_signatures) > 1: raise ValueError( "DP multi-concurrency requires all concurrent requests to " @@ -521,6 +518,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): if len(scheduler_output.scheduled_new_reqs) <= 1: return self.execute_request(scheduler_output) + parallel_config = getattr(self.od_config, "parallel_config", None) + dp_size = getattr(parallel_config, "data_parallel_size", 1) + if ( + dp_size > 1 + and getattr(self.od_config, "enable_distributed_layerwise_offload", False) + and getattr(self.od_config, "dlo_use_allgather", True) + ): + # DLO DP uses one independent request per DP replica. It is not a + # fused pipeline request batch, so models such as MiniMax-H3 do not + # need to advertise supports_request_batch=True. + return self.execute_request(scheduler_output) + result = self.collective_rpc( "execute_model_batch", args=(scheduler_output, self.od_config), @@ -599,7 +608,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): # so the D2H copy runs on a side stream in the worker background # thread while the default stream is free for the next forward. # pump routes the result back to a Future via rpc_id. - if not self.od_config.step_execution and method in ("execute_model", "execute_model_batch"): + multi_rank_reply = unique_reply_rank is None and exec_all_ranks + if ( + not self.od_config.step_execution + and method in ("execute_model", "execute_model_batch") + and not multi_rank_reply + ): rpc_id = self._next_rpc_id() rpc_request["rpc_id"] = rpc_id fut: concurrent.futures.Future = concurrent.futures.Future() diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index a8e23756f..bb589d7bb 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -221,9 +221,15 @@ class DenoiseProgressMixin: step_idx: int | None, timestep=None, scheduler=None, + normalized_timestep: float | None = None, ) -> None: set_forward_context_denoise_step_idx(step_idx) - if timestep is None or _forward_context is None: + if _forward_context is None: + return + if normalized_timestep is not None: + _forward_context.denoise_timestep = float(normalized_timestep) + return + if timestep is None: return scheduler = scheduler if scheduler is not None else getattr(self, "scheduler", None) ntt = getattr(getattr(scheduler, "config", None), "num_train_timesteps", None) diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 51c627975..14ab677a0 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -65,9 +65,9 @@ class LayerNorm(nn.LayerNorm, CustomOp): class RMSNorm(CustomOp): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None: super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) self.variance_epsilon = eps def _forward_fused(self, x: torch.Tensor) -> torch.Tensor: diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 6a045fca2..5d934866b 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -89,10 +89,11 @@ class RotaryEmbedding(CustomOp): of 1st half and 2nd half (GPT-NeoX style). """ - def __init__(self, is_neox_style: bool = False) -> None: + def __init__(self, is_neox_style: bool = False, half_head_dim: bool = True) -> None: super().__init__() self.is_neox_style = is_neox_style self.interleaved = not is_neox_style + self.half_head_dim = half_head_dim self.apply_rotary_emb_flash_attn = None self.has_mindie = False # ``find_spec("flash_attn")`` is True as long as *any* package publishes @@ -110,6 +111,21 @@ class RotaryEmbedding(CustomOp): if find_spec("mindiesd") is not None: self.has_mindie = True + def _prepare_half_head_dim_cos_sin( + self, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize shared cos/sin to the half-dim layout used by CUDA/native kernels.""" + if cos.dim() == 3: + # All batch elements share the same rotary position encoding. + cos = cos[0] + sin = sin[0] + if not self.half_head_dim: + cos = cos[..., ::2] if self.interleaved else cos[..., : cos.shape[-1] // 2] + sin = sin[..., ::2] if self.interleaved else sin[..., : sin.shape[-1] // 2] + return cos, sin + def forward_cuda( self, x: torch.Tensor, @@ -118,10 +134,7 @@ class RotaryEmbedding(CustomOp): ) -> torch.Tensor: from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb - if cos.dim() == 3: - # (B, S, D/2) -> (S, D/2) - cos = cos[0] - sin = sin[0] + cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin) x, squeezed = _ensure_batch_dim(x) output = apply_rotary_emb( @@ -141,10 +154,7 @@ class RotaryEmbedding(CustomOp): if self.apply_rotary_emb_flash_attn is None: return self.forward_cuda(x, cos, sin) - if cos.dim() == 3: - # (B, S, D/2) -> (S, D/2) - cos = cos[0] - sin = sin[0] + cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin) x, squeezed = _ensure_batch_dim(x) output = self.apply_rotary_emb_flash_attn( @@ -162,7 +172,7 @@ class RotaryEmbedding(CustomOp): sin: torch.Tensor, ) -> torch.Tensor: if self.has_mindie: - return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved) + return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved, self.half_head_dim) else: return self.forward_native(x, cos, sin) @@ -191,9 +201,7 @@ class RotaryEmbedding(CustomOp): # All batch elements share the same rotary position encoding. # Strip the batch dim so the underlying op broadcasts over the batch, # consistent with forward_cuda / forward_hip / apply_rotary_emb_mindiesd. - if cos.dim() == 3: - cos = cos[0] - sin = sin[0] + cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin) return apply_rotary_emb_torch( x, cos, diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 2fcf43b6e..9098fec2a 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -366,8 +366,8 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - # Default denoising components plus any a pipeline opts into via - # ``_lora_components`` (opt-in; other pipelines unchanged). + # Default denoising components, declared DiT components, and any a + # pipeline opts into via ``_lora_components``. # # NOTE: SDXL-style pipelines expose the denoiser as ``unet``. # Without scanning this component, adapters can load/activate while @@ -379,8 +379,10 @@ class DiffusionLoRAManager: "bagel", "unet", ) + declared_components = tuple(getattr(self.pipeline, "_dit_modules", ()) or ()) extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ()) - for component_name in (*default_components, *extra_components): + component_names = dict.fromkeys((*default_components, *declared_components, *extra_components)) + for component_name in component_names: if not hasattr(self.pipeline, component_name): continue component = getattr(self.pipeline, component_name) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index df7aa2523..d1a1ed15e 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -393,6 +393,13 @@ class DiffusersPipelineLoader: "DLO+AllGather active but model does not support mmap: loading weights via regular loader." ) logger.debug("Loading weights on %s ...", load_device) + if offload_after_quant: + marked = self._request_offload_after_quant(model) + if marked: + logger.info( + "Online quantization will return each of %d layers to CPU as it is quantized", + marked, + ) if load_format == "diffusers": cast(DiffusersAdapterPipeline, model).load_weights() else: @@ -406,6 +413,29 @@ class DiffusersPipelineLoader: self._apply_skip_softmax_calibration(model) return model.eval() + @staticmethod + def _request_offload_after_quant(model: nn.Module) -> int: + """Ask online-quant layers to return to host memory once quantized. + + The weights only visit the accelerator so the quant kernels can run on + them; ``load_model`` sends the model back to the host afterwards either + way. Without this the whole transformer accumulates on device until that + final move, which for MiniMax H3 is a ~43 GiB peak that no longer fits + beside a resident TP-sharded text encoder — even though layer-wise + offload means none of it is supposed to be resident at inference time. + + Only quant methods that advertise ``supports_offload_after_quant`` are + asked, since the implementation has to know when a layer is finished. + Deferring materialization to the ``meta`` device does not imply that. + """ + marked = 0 + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if getattr(quant_method, "supports_offload_after_quant", False): + quant_method.enable_offload_after_quant() + marked += 1 + return marked + @staticmethod def _has_online_quant(model: nn.Module) -> bool: """Whether any layer uses an online-quant method that defers weight @@ -460,6 +490,13 @@ class DiffusersPipelineLoader: if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): continue + # Layers finished during loading would only be staged onto the target + # device for a process call that immediately returns. That round trip + # is wasted work in general, and undoes the point of the offload for + # layers that already went back to the host. + if getattr(module, "_already_called_process_weights_after_loading", False): + continue + if has_online_quant: # Online quant may leave straggler params on the ``meta`` device. # Move only real (non-meta) params onto the target device for diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index 22c9a2737..6184478b8 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -41,6 +41,16 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=9, supports_mixed_reference_inputs=True, + # H3 represents alignment padding as a second packed sequence. The + # packed TRTLLM backend consumes cu_seqlens and isolates that padding. + attention_mask_free=True, + ), + # The modular alias is served by MiniMaxH3Pipeline and has the same + # Ref2VA request contract. Keep admission limits in sync with it. + "MiniMaxH3ModularPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=9, + supports_mixed_reference_inputs=True, ), "WanPipeline": DiffusionModelMetadata(attention_mask_free=True), "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True), diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index 45f1e653e..2369a3b05 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -14,6 +14,9 @@ from typing import Any import torch +from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout +from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx + from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, @@ -88,6 +91,13 @@ class MiniMaxH3DenoiseBranch: "max_seqlen_q": text_len, }, } + # Where the video segment sits in the packed sequence. Resolved to plain + # ints here so the attention layers never sync on it per step. + grid = packed["latent_grid"].tolist() + self.static_kwargs["video_token_layout"] = VideoTokenLayout( + prefix_len=int(packed["video_row_start"]), + latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])), + ) def forward_kwargs( self, @@ -140,7 +150,8 @@ def minimax_h3_denoise_loop( device: torch.device, imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP, audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, - on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, + on_step_end: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, + on_step_start: Callable[[int, float, float], None] | None = None, step_profiler: Callable[[int], AbstractContextManager] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Run the full denoise loop; returns final (video_rows, audio_rows). @@ -193,8 +204,14 @@ def minimax_h3_denoise_loop( for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: + # Publish the step index so step-gated attention features (e.g. the + # dense warmup of RAINFUSION_ATTN) can see where we are. + set_forward_context_denoise_step_idx(step) s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] + if on_step_start is not None: + # Attention gates use the scheduler-style descending timestep. + on_step_start(step, s_v, s_a) t_v, t_a = 1.0 - s_v, 1.0 - s_a imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) @@ -247,9 +264,10 @@ def minimax_h3_denoise_loop( audio_rows[audio_update] = new_audio if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor # per-step audio ref reset - if on_step is not None: - on_step(step, video_rows, audio_rows) + if on_step_end is not None: + on_step_end(step, video_rows, audio_rows) + set_forward_context_denoise_step_idx(None) return video_rows, audio_rows diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 438ff80a2..f811eeecb 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -1029,16 +1029,84 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): def load_to_device(self) -> None: if not self.is_loaded: return + if getattr(self, "_omni_layerwise_enabled", False): + for name, child in self.vision.named_children(): + if name != "blocks": + child.to(self.device_target) + for name, child in self.text_model.named_children(): + if name != "layers": + child.to(self.device_target) + return self.vision.to(self.device_target) self.text_model.to(self.device_target) def offload_to_cpu(self) -> None: if not self.is_loaded: return + if getattr(self, "_omni_layerwise_enabled", False): + for hook in self._omni_layerwise_hooks: + hook.offload_layer() + for name, child in self.vision.named_children(): + if name != "blocks": + child.to("cpu") + for name, child in self.text_model.named_children(): + if name != "layers": + child.to("cpu") + torch.accelerator.empty_cache() + return self.vision.to("cpu") self.text_model.to("cpu") torch.accelerator.empty_cache() + def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None: + """Stream the TP-local Qwen vision/text blocks for low-HBM serving. + + The encoder has its own TP process group, so these blocks must remain + rank-local. Reusing the DiT AllGather group would concatenate + different TP shards and corrupt the encoder weights. + """ + if not self.is_loaded or getattr(self, "_omni_layerwise_enabled", False): + return + + from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook + from vllm_omni.platforms import current_omni_platform + + self._omni_layerwise_hooks = [] + self._omni_layerwise_block_groups = [] + copy_stream = current_omni_platform.Stream() + for blocks in (self.vision.blocks, self.text_model.layers): + if len(blocks) <= 1: + continue + last_hook = apply_block_hook( + blocks[-1], + blocks[0], + self.device_target, + copy_stream, + pin_memory, + ) + hooks = [last_hook] + for index, block in enumerate(blocks[:-1]): + hooks.append( + apply_block_hook( + block, + blocks[index + 1], + self.device_target, + copy_stream, + pin_memory, + ) + ) + for index, hook in enumerate(hooks): + hook._prev_hook = hooks[index - 1] + self._omni_layerwise_hooks.extend(hooks) + self._omni_layerwise_block_groups.append(blocks) + + self._omni_layerwise_enabled = bool(self._omni_layerwise_hooks) + logger.info( + "MiniMax H3 encoder layerwise offload enabled on %d blocks across %d stacks", + sum(len(blocks) for blocks in self._omni_layerwise_block_groups), + len(self._omni_layerwise_block_groups), + ) + def _encode( self, input_ids: torch.Tensor, diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index fd5ca67e9..610c20d19 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -26,13 +26,15 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.diffusion.layers.rope import RotaryEmbedding if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -132,6 +134,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( "img_pos_for_infer_output_info", "packed_seq_params", "refiner_packed_seq_params", + "video_token_layout", } ) @@ -168,16 +171,11 @@ def _reorder_grouped_qkv_to_qkv( ) -def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> nn.RMSNorm: +def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> RMSNorm: # RMSNorm uses fp32 accumulation with bf16 inputs and outputs. # torch.nn.RMSNorm upcasts reduced-precision inputs for the variance # reduction, matching that accumulation semantic. - return nn.RMSNorm(size, eps=eps, dtype=dtype) - - -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - x1, x2 = torch.chunk(x, 2, dim=-1) - return torch.cat((-x2, x1), dim=-1) + return RMSNorm(size, eps=eps, dtype=dtype) def _modulate_scale_shift( @@ -230,20 +228,6 @@ class MiniMaxH3Rope(nn.Module): return torch.cat((half, half), dim=-1) # [S, 96] -def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: - """Rotate the first rot_dim head dims; pass the rest through. - - x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin - are cast to the activation dtype before the elementwise math. - """ - rot_dim = freqs.shape[-1] - x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:] - cos = torch.cos(freqs).to(x.dtype).unsqueeze(1) # [T, 1, rot_dim] - sin = torch.sin(freqs).to(x.dtype).unsqueeze(1) - x_rot = (x_rot * cos) + (_rotate_half(x_rot) * sin) - return torch.cat((x_rot, x_pass), dim=-1) - - class MiniMaxH3TimeEmbedder(nn.Module): def __init__( self, @@ -327,6 +311,8 @@ class MiniMaxH3Attention(nn.Module): quant_config: QuantizationConfig | None, *, prefix: str, + role: str = "self", + role_category: str | None = None, skip_sequence_parallel: bool = False, ) -> None: super().__init__() @@ -349,6 +335,7 @@ class MiniMaxH3Attention(nn.Module): self.num_kv_heads = self.qkv_proj.num_kv_heads self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps) self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps) + self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False) self.out_proj = RowParallelLinear( inner_dim, arch.hidden_size, @@ -364,9 +351,27 @@ class MiniMaxH3Attention(nn.Module): head_size=self.head_dim, softmax_scale=self.softmax_scale, causal=False, + # Packed rows reach the impl as [B, S, N, D]. + qkv_layout="BSND", + role=role, + role_category=role_category, skip_sequence_parallel=skip_sequence_parallel, + prefix=prefix, ) + def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Rotate the first rot_dim head dims; pass the rest through. + + x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin + are cast to the activation dtype before the elementwise math. + """ + rot_dim = freqs.shape[-1] + x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:] + cos = torch.cos(freqs).to(x.dtype) # [T, rot_dim] + sin = torch.sin(freqs).to(x.dtype) + x_rot = self.rope(x_rot, cos, sin) + return torch.cat((x_rot, x_pass), dim=-1) + @torch.compiler.disable def _run_packed_attention( self, @@ -376,18 +381,31 @@ class MiniMaxH3Attention(nn.Module): *, cu_seqlens: torch.Tensor, max_seqlen: int, + packed_total: int, + video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """Run packed attention as a small eager island. - The scalar packed-layout metadata and the CuTe FlashAttention-4 DSL - are intentionally opaque to Dynamo. Keeping this boundary narrow lets - regional compile fuse projections, norms, RoPE, and the surrounding - DiT block without repeatedly graph-breaking inside the FA4 compiler. + The scalar packed-layout metadata and backend-specific attention + kernels are intentionally opaque to Dynamo. Keeping this boundary + narrow lets regional compile fuse projections, norms, RoPE, and the + surrounding DiT block without repeated graph breaks. """ - used = int(cu_seqlens[1].item()) - packed_total = int(cu_seqlens[-1].item()) + # max_seqlen is already the first (real) packed document length. Do + # not read the CUDA cu_seqlens scalars here: this function runs once + # per layer and .item() would serialize every attention launch. + if not 0 < max_seqlen <= packed_total: + raise ValueError( + f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}" + ) + used = min(max_seqlen, packed_total) attn_mask = None - if used < packed_total: + # Ring attention can dispatch to a different implementation from the + # configured backend, so this no-mask fast path is local-only. + prefix_slice = ( + not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing + ) + if used < packed_total and not prefix_slice: attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, @@ -396,7 +414,9 @@ class MiniMaxH3Attention(nn.Module): "cu_seqlens_k": cu_seqlens, "max_seqlen_q": max_seqlen, "max_seqlen_k": max_seqlen, + "valid_kv_length": used, }, + video_layout=video_layout, ) return self.attention( q.unsqueeze(0), @@ -412,7 +432,9 @@ class MiniMaxH3Attention(nn.Module): rope_freqs: torch.Tensor | None, cu_seqlens: torch.Tensor, max_seqlen: int, + packed_total: int | None = None, sp_seq_lens: list[int] | None = None, + video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """x: [T, hidden] packed thd rows -> [T, hidden]. @@ -436,8 +458,8 @@ class MiniMaxH3Attention(nn.Module): q = self.q_norm(q) k = self.k_norm(k) if rope_freqs is not None: - q = _apply_rope(q, rope_freqs) - k = _apply_rope(k, rope_freqs) + q = self._apply_rope(q, rope_freqs) + k = self._apply_rope(k, rope_freqs) # The packed layout uses a second document for alignment padding. # Local/Ulysses backends unpad it, while Ring keeps aligned rows for @@ -448,6 +470,11 @@ class MiniMaxH3Attention(nn.Module): v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + # Before Ulysses, q contains only this rank's row shard. The + # backend receives the global sequence after all-to-all, so carry + # its Python length explicitly instead of inferring it from q. + packed_total=packed_total if packed_total is not None else q.shape[0], + video_layout=video_layout, ) out = out.reshape(total, self.num_heads * self.head_dim) out, _ = self.out_proj(out) @@ -558,6 +585,8 @@ class MiniMaxH3TokenRefinerBlock(nn.Module): arch, quant_config, prefix=f"{prefix}.attn", + role="minimax_h3.token_refiner", + role_category="self", skip_sequence_parallel=True, ) self.mlp = MiniMaxH3MLP( @@ -627,6 +656,8 @@ class MiniMaxH3DiTBlock(nn.Module): super().__init__() self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps) + # The prefix also carries the block index that block-sparse attention + # backends match against their skip_layers selector. self.attn = MiniMaxH3Attention( arch, quant_config, @@ -655,7 +686,9 @@ class MiniMaxH3DiTBlock(nn.Module): rope_freqs: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, + packed_total: int, sp_seq_lens: list[int] | None = None, + video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T] (= inverse_indices * modality_num + token_tags.clamp(min=0)). @@ -681,7 +714,9 @@ class MiniMaxH3DiTBlock(nn.Module): rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + packed_total=packed_total, sp_seq_lens=sp_seq_lens, + video_layout=video_layout, ) x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) @@ -814,6 +849,10 @@ class MiniMaxH3DiTModel(nn.Module): }, "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2), } + # The checkpoint already stores qkv and the MLP gate/up as single tensors + # (see the reordering in load_weights), so there are no unfused names for + # quantization or LoRA to map onto. Address the fused layers directly, e.g. + # ignored_layers=["blocks.0.attn.qkv_proj"]. packed_modules_mapping = {} def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None: @@ -1080,6 +1119,7 @@ class MiniMaxH3DiTModel(nn.Module): refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) + video_layout = kwargs.get("video_token_layout") if x.dim() != 3 or x.shape[0] != 1: raise ValueError(f"x must be [1, S, C], got {list(x.shape)}") @@ -1127,6 +1167,8 @@ class MiniMaxH3DiTModel(nn.Module): rope_freqs=block_rope, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + packed_total=seq_len, + video_layout=video_layout, ) hidden = self.sp_gather(hidden) diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py index 37353b459..a57eae700 100644 --- a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py +++ b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py @@ -236,6 +236,9 @@ def minimax_h3_packed_sequence( "token_tags": token_tags, "cu_seqlens": cu, "document_id": doc, + # Extent of the video segment within the packed sequence. + "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64), + "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64), } @@ -554,6 +557,9 @@ def minimax_h3_packed_sequence_ref2va_blocks( "token_tags": token_tags, "cu_seqlens": cu, "document_id": doc, + # Extent of the video segment within the packed sequence. + "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64), + "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64), } diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 9ab54765a..b11b4e320 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -8,6 +8,8 @@ import math import os import tempfile from collections.abc import Iterable, Mapping, Sequence +from contextlib import contextmanager +from itertools import groupby from pathlib import Path from typing import Any, ClassVar @@ -26,6 +28,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( init_world_group, ) from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.forward_context import DenoiseProgressMixin from vllm_omni.diffusion.model_loader.diffusers_loader import ( DiffusersPipelineLoader, ) @@ -36,11 +39,15 @@ from vllm_omni.diffusion.models.interface import ( SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.offloader import OffloadPlan from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.errors import OmniClientError +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) from vllm_omni.platforms import current_omni_platform from .condition_noise import ( @@ -104,6 +111,57 @@ MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0 MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0 +MINIMAX_H3_DOWNLOAD_PATTERNS = [ + "FL2VA/**", + "Ref2VA/model_index.json", + "Ref2VA/transformer/**", +] +MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = { + "fl2va": ["FL2VA/**"], + "ref2va": ["Ref2VA/**"], +} + + +def _minimax_h3_partition_for_task( + task_type: str | None, + model: str | None = None, +) -> str: + task = str(task_type or "auto").lower() + if task == "auto" and model is not None: + path = Path(model) + if path.is_dir() and path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file(): + return path.name.lower() + if task in {"auto", "combined"}: + return "combined" + if task in {"t2va", "fl2va"}: + return "fl2va" + if task == "ref2va": + return "ref2va" + raise ValueError(f"MiniMax-H3 task_type must be one of auto, t2va, fl2va, or ref2va; got {task_type!r}") + + +def _resolve_minimax_h3_model_root( + model: str, + revision: str | None, + partition: str, +) -> Path: + path = Path(model) + if path.is_dir(): + if path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file(): + return path.parent + return path + allow_patterns = ( + MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition] + ) + return Path( + download_weights_from_hf_specific( + model_name_or_path=model, + cache_dir=None, + allow_patterns=allow_patterns, + revision=revision, + require_all=True, + ) + ) def _resolve_component_quant_config(quant_config, component: str): @@ -444,6 +502,7 @@ class _SingleRankEncoderGroup: class MiniMaxH3Pipeline( nn.Module, + DenoiseProgressMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportImageInput, @@ -453,9 +512,22 @@ class MiniMaxH3Pipeline( ): """CFG-distilled joint video/audio generation for MiniMax H3.""" - _dit_modules: ClassVar[list[str]] = ["transformer"] + _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"] + _offload_plan: ClassVar[OffloadPlan] = OffloadPlan( + offload_submodules={"token_refiner": "blocks"}, + resident_dit_paths=frozenset({"transformer"}), + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), + ) + # H3's regular loader performs checkpoint-layout conversions (grouped QKV + # and fused MLP). Keep it on that path until mmap can run the same loader + # callbacks for every affected parameter. + _supports_mmap_loading: ClassVar[bool] = False + # TODO(offload): Re-enable after the generic rank-local mmap path can run + # this model's grouped-QKV reorder and fused-MLP packing before TP sharding. + # Do not bypass the regular loader until that equivalence is tested. _PROFILER_TARGETS: ClassVar[list[str]] = [ "_prepare_reference_videos", "encode_prompt", @@ -479,24 +551,64 @@ class MiniMaxH3Pipeline( if int(self.parallel_config.cfg_parallel_size) != 1: raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1") self.device = get_local_device() - model_path = str(od_config.model) - model_index = json.loads((Path(model_path) / "model_index.json").read_text(encoding="utf-8")) + self.partition = _minimax_h3_partition_for_task( + getattr(od_config, "task_type", None), + str(od_config.model), + ) + model_root = _resolve_minimax_h3_model_root( + str(od_config.model), + od_config.revision, + self.partition, + ) + model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA") + model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8")) release = model_index.get("_minimax_h3") or {} - self.partition = str(release.get("partition", "")) - self.supported_tasks = frozenset(release.get("tasks") or ()) + partition = str(release.get("partition", "")).lower() + expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va" + if partition != expected_partition: + raise ValueError(f"invalid MiniMax-H3 {expected_partition} partition at {model_path}") + + supported_tasks = {str(task).lower() for task in release.get("tasks", [])} + if not supported_tasks: + supported_tasks = {"ref2va"} if partition == "ref2va" else {"t2va", "fl2va"} + ref2va_model_path = None + if self.partition == "combined": + ref2va_model_path = model_root / "Ref2VA" + ref2va_index_path = ref2va_model_path / "model_index.json" + if not ref2va_index_path.is_file(): + raise ValueError(f"Ref2VA partition not found at {ref2va_model_path}") + ref2va_index = json.loads(ref2va_index_path.read_text(encoding="utf-8")) + ref2va_release = ref2va_index.get("_minimax_h3") or {} + if str(ref2va_release.get("partition", "")).lower() != "ref2va": + raise ValueError(f"invalid MiniMax-H3 ref2va partition at {ref2va_model_path}") + supported_tasks.update(str(task).lower() for task in ref2va_release.get("tasks", ["ref2va"])) + + self.supported_tasks = frozenset(supported_tasks) shifts = release.get("sigma_shift_scales") or {} self.default_video_shift = float(shifts.get("video", 12.0)) self.default_audio_shift = float(shifts.get("audio", 3.0)) self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, + model_or_path=str(model_path), subfolder="transformer", revision=od_config.revision, prefix="transformer.", fall_back_to_pt=False, ) ] + self._dit_modules = ["transformer"] + if ref2va_model_path is not None: + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=str(ref2va_model_path), + subfolder="transformer", + revision=od_config.revision, + prefix="transformers_ref.", + fall_back_to_pt=False, + ) + ) + self._dit_modules.append("transformers_ref") transformer_quant_config = _resolve_component_quant_config( od_config.quantization_config, "transformer", @@ -505,14 +617,19 @@ class MiniMaxH3Pipeline( od_config, quant_config=transformer_quant_config, ) + if ref2va_model_path is not None: + self.transformers_ref = MiniMaxH3DiTModel( + od_config, + quant_config=transformer_quant_config, + ) self.tokenizer = Qwen2TokenizerFast.from_pretrained( - model_path, + str(model_path), subfolder="tokenizer", local_files_only=os.path.isdir(model_path), ) self.processor = Qwen3VLProcessor.from_pretrained( - model_path, + str(model_path), subfolder="processor", local_files_only=os.path.isdir(model_path), ) @@ -542,13 +659,19 @@ class MiniMaxH3Pipeline( load_model=rank < text_encoder_tp_size, encoder_group=self.text_encoder_group, ) + stage_components = bool( + od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) + ) + component_load_device = torch.device("cpu") if stage_components else self.device self.video_vae = MiniMaxH3VideoVAE( os.path.join(model_path, "video_vae"), device=self.device, + load_device=component_load_device, ) self.audio_vae = MiniMaxH3AudioVAE( os.path.join(model_path, "audio_vae"), device=self.device, + load_device=component_load_device, ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae @@ -561,16 +684,23 @@ class MiniMaxH3Pipeline( self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - prefix = "transformer." + def source_prefix(item: tuple[str, torch.Tensor]) -> str: + name, _ = item + prefix = name.partition(".")[0] + "." + if prefix in {"transformer.", "transformers_ref."}: + return prefix + raise ValueError(f"unexpected MiniMax-H3 weight {name!r}") - def transformer_weights(): - for name, tensor in weights: - if name.startswith(prefix): - yield name[len(prefix) :], tensor - - loaded = self.transformer.load_weights(transformer_weights()) - self.transformer.post_load_weights() - loaded_with_prefix = {prefix + name for name in loaded} + loaded_with_prefix: set[str] = set() + loaded_prefixes: set[str] = set() + for prefix, grouped_weights in groupby(weights, key=source_prefix): + if prefix in loaded_prefixes: + raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous") + loaded_prefixes.add(prefix) + transformer = getattr(self, prefix.removesuffix(".")) + loaded = transformer.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights) + transformer.post_load_weights() + loaded_with_prefix.update(prefix + name for name in loaded) # The text encoder and both VAEs load eagerly in ``__init__`` rather # than through ``weights_sources``. Record them for the runner's strict # missing-parameter check. @@ -579,14 +709,23 @@ class MiniMaxH3Pipeline( loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) return loaded_with_prefix + def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel: + if task == "ref2va" and hasattr(self, "transformers_ref"): + return self.transformers_ref + return self.transformer + def _resolve_task( self, requested: str | None, multi_modal_data: dict[str, Any], ) -> str: if requested is None: + # A Ref2VA-only startup has no FL2VA transformer; preserve its + # historical implicit default even for image-only references. if self.partition == "ref2va": requested = "ref2va" + elif multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None: + requested = "ref2va" elif multi_modal_data.get("image") is not None: requested = "fl2va" else: @@ -935,12 +1074,16 @@ class MiniMaxH3Pipeline( input_ids: torch.Tensor, vision_kwargs: dict[str, torch.Tensor], ) -> torch.Tensor: - if self.od_config.enable_cpu_offload: + if self.od_config.enable_cpu_offload and not getattr( + self.od_config, "enable_distributed_layerwise_offload", False + ): # Invoke nn.Module.__call__ so the generic model-level offloader # swaps the resident DiT and encoder. return self.text_encoder(input_ids, **vision_kwargs) - if self.od_config.enable_layerwise_offload: + if self.od_config.enable_layerwise_offload or getattr( + self.od_config, "enable_distributed_layerwise_offload", False + ): # Layerwise DiT offload already provides the low-residency encoder # phase used by the checkpoint reference. self.text_encoder.load_to_device() @@ -955,12 +1098,33 @@ class MiniMaxH3Pipeline( self.text_encoder.load_to_device() return self.text_encoder.encode_ids(input_ids, **vision_kwargs) + def _uses_manual_component_offload(self) -> bool: + od_config = getattr(self, "od_config", None) + return bool( + getattr(od_config, "enable_layerwise_offload", False) + or getattr(od_config, "enable_distributed_layerwise_offload", False) + ) + + @contextmanager + def _component_on_device(self, component: nn.Module): + staged = self._uses_manual_component_offload() + if staged: + component.load_to_device() + try: + yield + finally: + if staged: + component.offload_to_cpu() + def _encode_visual_condition( self, image: Image.Image, ) -> torch.Tensor: _, rank, _ = _dit_rank_world() - rows = self.video_vae.encode_image(image) if rank == 0 else None + rows = None + if rank == 0: + with self._component_on_device(self.video_vae): + rows = self.video_vae.encode_image(image) return _broadcast_tensor( rows, dtype=torch.float32, @@ -996,7 +1160,8 @@ class MiniMaxH3Pipeline( rows = None audio_t = 0 if rank == 0: - rows, audio_t = self.audio_vae.encode_waveform(*audio) + with self._component_on_device(self.audio_vae): + rows, audio_t = self.audio_vae.encode_waveform(*audio) audio_t_tensor = torch.tensor( [audio_t], dtype=torch.long, @@ -1042,6 +1207,15 @@ class MiniMaxH3Pipeline( prepared_videos: list[dict[str, Any]] | None, *, count: int, + ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: + with self._component_on_device(self.video_vae): + return self._encode_video_conditions_resident(prepared_videos, count=count) + + def _encode_video_conditions_resident( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + count: int, ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: group, rank, world_size = _dit_rank_world() distributed_encode = self.video_vae.is_distributed_enabled() @@ -1089,6 +1263,15 @@ class MiniMaxH3Pipeline( prepared_videos: list[dict[str, Any]] | None, *, has_audio: list[bool], + ) -> tuple[torch.Tensor | None, list[int]]: + with self._component_on_device(self.audio_vae): + return self._encode_video_audio_conditions_resident(prepared_videos, has_audio=has_audio) + + def _encode_video_audio_conditions_resident( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + has_audio: list[bool], ) -> tuple[torch.Tensor | None, list[int]]: _, rank, _ = _dit_rank_world() count = sum(has_audio) @@ -1156,6 +1339,17 @@ class MiniMaxH3Pipeline( ) return video_rows, audio_rows + @contextmanager + def _resident_dit_layers_on_device(self, *, enabled: bool = True): + controller = getattr(self, "_dlo_residency_controller", None) + if controller is not None and enabled: + controller.load_resident_layers() + try: + yield + finally: + if controller is not None and enabled: + controller.offload_resident_layers() + def diffuse( self, *, @@ -1277,21 +1471,31 @@ class MiniMaxH3Pipeline( num_steps=num_steps, shift_scale=audio_shift, ) - with self.progress_bar(total=len(video_sigmas) - 1) as progress: - video_rows, audio_rows = minimax_h3_denoise_loop( - model=self.transformer, - positive=branch, - initial_video_rows=initial_video, - initial_audio_rows=initial_audio, - keyframe_cond_rows=visual_anchor, - audio_ref_rows=audio_anchor, - sigmas_video=video_sigmas, - sigmas_audio=audio_sigmas, - device=self.device, - imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), - audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), - on_step=lambda step, video, audio: progress.update(), - ) + transformer = self._transformer_for_task(task) + # The static DLO plan keeps leading blocks resident only for the + # primary ``transformer``. In combined mode ``transformers_ref`` is + # fully streamed, so a Ref2VA request must not stage the inactive + # FL2VA transformer resident blocks. + with self._resident_dit_layers_on_device(enabled=transformer is self.transformer): + with self.progress_bar(total=len(video_sigmas) - 1) as progress: + video_rows, audio_rows = minimax_h3_denoise_loop( + model=transformer, + positive=branch, + initial_video_rows=initial_video, + initial_audio_rows=initial_audio, + keyframe_cond_rows=visual_anchor, + audio_ref_rows=audio_anchor, + sigmas_video=video_sigmas, + sigmas_audio=audio_sigmas, + device=self.device, + imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), + audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + on_step_start=lambda step, video_sigma, audio_sigma: self.record_denoise_step( + step, + normalized_timestep=video_sigma, + ), + on_step_end=lambda step, video, audio: progress.update(), + ) target_video = video_rows[branch.update_mask_dev] video_latent = minimax_h3_unpatchify_video_tokens( @@ -1320,14 +1524,16 @@ class MiniMaxH3Pipeline( height: int, width: int, ) -> tuple[torch.Tensor, torch.Tensor]: - with current_omni_platform.create_autocast_context( - device_type=self.device.type, - dtype=torch.float16, - enabled=True, - ): - video = self.video_vae.decode_latent(video_latent) + with self._component_on_device(self.video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + video = self.video_vae.decode_latent(video_latent) video = video[..., :height, :width].contiguous() - audio = self.audio_vae.decode_latent(audio_latent) + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) return video, audio @torch.no_grad() diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 8f2152dc1..e28450c90 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -19,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor DistributedVaeMixin, ) from vllm_omni.diffusion.distributed.parallel_state import get_dit_group +from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from .packed_tokens import minimax_h3_patchify_video_latent @@ -50,7 +51,14 @@ def _load_remote_component( class_reference, component_path, ) - return component_cls.from_pretrained(component_path) + # Build on the host regardless of the ambient default device. Online + # quantization wraps pipeline construction in a `with torch.device()` + # block for the DiT's quantized linears, and the checkpoint's own VAE code + # builds constants with ops that have no accelerator kernel (BigVGAN's + # anti-aliasing filters call torch.kaiser_window). Callers place the module + # explicitly right after this returns, so nothing depends on the context. + with torch.device("cpu"): + return component_cls.from_pretrained(component_path) class _AudioVAEDeterminismContext(AbstractContextManager): @@ -102,8 +110,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): component_path: str, *, device: torch.device, + load_device: torch.device | None = None, ) -> None: super().__init__() + self._device_target = device self.config_dict = _load_component_config(component_path) self.remote = _load_remote_component( component_path, @@ -112,11 +122,33 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): # Match the reference loader contract: video VAE weights stay FP32. # Keyframe encoding is numerically sensitive to first casting the # checkpoint through FP16; decode still runs under FP16 autocast. - self.remote.eval().to(device=device, dtype=torch.float32) + initial_device = load_device or device + self.remote.eval().to(device=initial_device, dtype=torch.float32) + self._stager = None + if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): + self._stager = PinnedModuleStager( + self.remote, + device, + pin_memory=True, + ) self.model = self.remote.model self.use_tiling = True self.use_slicing = False self.parallel_size = 1 + self.device_module = torch.get_device_module() + + def load_to_device(self) -> None: + if self._stager is not None: + self._stager.load() + else: + self.remote.to(self._device_target) + + def offload_to_cpu(self) -> None: + if self._stager is not None: + self._stager.offload() + else: + self.remote.to("cpu") + torch.accelerator.empty_cache() def set_parallel_size( self, @@ -166,13 +198,13 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): previous_dtype = parameter.dtype if previous_dtype != torch.float32: self.to(torch.float32) - devices = [parameter.device] if parameter.device.type == "cuda" else [] + devices = [parameter.device] if parameter.device.type != "cpu" else [] try: - with torch.random.fork_rng(devices=devices): + with torch.random.fork_rng(devices=devices, device_type=parameter.device.type): torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) for device in devices: - with torch.cuda.device(device): - torch.cuda.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) + with self.device_module.device(device): + self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) latent = self.model.encode_images( image, use_fp16_latent=True, @@ -210,13 +242,13 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): previous_dtype = parameter.dtype if previous_dtype != torch.float32: self.to(torch.float32) - devices = [parameter.device] if parameter.device.type == "cuda" else [] + devices = [parameter.device] if parameter.device.type != "cpu" else [] try: - with torch.random.fork_rng(devices=devices): + with torch.random.fork_rng(devices=devices, device_type=parameter.device.type): torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) for device in devices: - with torch.cuda.device(device): - torch.cuda.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) + with self.device_module.device(device): + self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) latent = self.model.encode_videos( frames, use_fp16_latent=True, @@ -276,8 +308,10 @@ class MiniMaxH3AudioVAE(nn.Module): component_path: str, *, device: torch.device, + load_device: torch.device | None = None, ) -> None: super().__init__() + self._device_target = device self.config_dict = _load_component_config(component_path) self.remote = _load_remote_component( component_path, @@ -285,10 +319,31 @@ class MiniMaxH3AudioVAE(nn.Module): ) # The checkpoint's audio VAE contract is FP32 for both reference # encoding and waveform decoding. - self.remote.eval().to(device=device, dtype=torch.float32) + initial_device = load_device or device + self.remote.eval().to(device=initial_device, dtype=torch.float32) + self._stager = None + if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): + self._stager = PinnedModuleStager( + self.remote, + device, + pin_memory=True, + ) self.model = self.remote.model self.sample_rate = int(self.config_dict["sample_rate"]) + def load_to_device(self) -> None: + if self._stager is not None: + self._stager.load() + else: + self.remote.to(self._device_target) + + def offload_to_cpu(self) -> None: + if self._stager is not None: + self._stager.offload() + else: + self.remote.to("cpu") + torch.accelerator.empty_cache() + @torch.inference_mode() def encode_waveform( self, diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 807fad7ed..9ff464494 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -16,6 +16,7 @@ from .distributed_layerwise_backend import ( remove_distributed_block_hook, ) from .layerwise_backend import LayerWiseOffloadBackend +from .module_residency import PinnedModuleStager from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading from .sequential_backend import ( ModelLevelOffloadBackend, @@ -41,6 +42,7 @@ __all__ = [ "DistributedLayerwiseOffloadBackend", "DistributedLayerwiseOffloadHook", "ModelLevelOffloadBackend", + "PinnedModuleStager", "apply_sequential_offload", "remove_sequential_offload", "apply_distributed_block_hook", diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 20b2f5e14..2db80923d 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -27,7 +27,10 @@ class OffloadConfig: pin_cpu_memory: bool = True use_hsdp: bool = False dp_size: int = 1 # derived from parallel_config, not user-configurable - dlo_use_allgather: bool = True # if True: shard + AllGather; if False: full weights H2D only + # True: add DP sharding + AllGather. False: stream the standard loader's + # rank-local tensors (including TP-local shards) with H2D only. + dlo_use_allgather: bool = True + dlo_resident_layers: int = 0 # leading DiT layers kept on device model_path: str | None = None # checkpoint path for mmap weight loading @classmethod @@ -94,17 +97,26 @@ class OffloadConfig: else: strategy = OffloadStrategy.NONE - # When dlo_use_allgather=False, each rank loads full weights (no sharding, - # no AllGather). This avoids AllGather synchronization requirements - # (concurrent requests, dummy run skip) at the cost of N× CPU memory. + # With dlo_use_allgather=False, do not add another DP shard. Each rank + # streams the tensors produced by the standard loader, which may + # already be TP-local shards. This avoids AllGather synchronization + # requirements (concurrent requests, dummy run skip). dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) + dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0)) + if dlo_resident_layers < 0: + raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}") + if dlo_resident_layers and dlo_use_allgather: + raise ValueError( + "dlo_resident_layers currently requires --dlo-no-use-allgather so " + "resident blocks use weights prepared by the standard TP-aware loader" + ) # If dlo_use_allgather=False, force dp_size=1 (each rank independent) if enable_distributed_layerwise_offload and not dlo_use_allgather: dp_size = 1 logger.info( "Distributed layerwise offload: dlo_use_allgather=False, " - "each rank loads full weights (no shard, no AllGather)" + "streaming standard-loader rank-local weights (no DP shard or AllGather)" ) # HSDP already shards parameters into DTensors. Running distributed @@ -114,8 +126,8 @@ class OffloadConfig: raise ValueError( "Distributed layerwise offload with AllGather is incompatible with " "HSDP: HSDP parameters are already sharded DTensors, and the offloader " - "would double-shard them. Use --dlo-no-use-allgather (full weights per " - "rank) or disable HSDP." + "would double-shard them. Use --dlo-no-use-allgather (standard-loader " + "rank-local weights) or disable HSDP." ) return cls( @@ -124,6 +136,7 @@ class OffloadConfig: use_hsdp=use_hsdp, dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, + dlo_resident_layers=dlo_resident_layers, model_path=getattr(od_config, "model", None), ) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 797209781..7da7c8d37 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -219,6 +219,15 @@ class DistributedLayerwiseOffloadHook(ModelHook): weights_with_local = [] for name, t in name2weights.items(): local_t = t.to_local() if hasattr(t, "to_local") else t + mmap_transform = getattr(t, "mmap_weight_transform", None) + if callable(mmap_transform) and getattr(t, "mmap_weight_transform_pending", False): + # Some checkpoints use a layout that is converted by the + # regular weight loader (for example MiniMax-H3 grouped + # QKV). Apply that conversion one block at a time while + # copying the rank-local CPU shard. Keeping the raw + # parameter as an mmap view avoids a private full-model + # copy in every worker. + local_t = mmap_transform(local_t) weights_with_local.append((name, t, local_t)) total_numel = sum(local.numel() for _, _, local in weights_with_local) @@ -330,7 +339,15 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.comm_stream.wait_stream(self.copy_stream) with current_omni_platform.stream(self.comm_stream): for dtype, local_shard in gpu_shards.items(): - gw = gpu_weights[dtype] + # Slice the shared (max-sized) output buffer down to this + # block's actual AllGather output size. The buffers are + # sized to the *largest* block across all groups, so for any + # smaller block the full buffer would violate the + # all_gather_into_tensor contract + # (output.numel() == world_size * input.numel()). + # Repoint offsets are relative to the block's flattened + # buffer, so a prefix slice is safe. + gw = gpu_weights[dtype][: self._ag_output_sizes[dtype]] torch.distributed.all_gather_into_tensor( gw, local_shard, @@ -477,6 +494,111 @@ def remove_distributed_block_hook(module: nn.Module) -> None: logger.debug("Removed distributed offload hook from %s", module.__class__.__name__) +class PinnedResidentLayerGroup: + """Keep selected layers in pinned host memory between requests. + + + TODO(offload): Extract this alongside PinnedModuleStager after the + distributed shard-and-pin operation becomes a shared storage primitive. + It currently remains here because it depends on DLO's local-shard layout. + Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a + pinned CPU master copy and never copies generated device weights back to + host. Entering the denoise stage performs one asynchronous H2D pass; + leaving it only restores zero-sized placeholders and releases the device + buffers. This lets the following VAE stage reuse the same HBM. + + The resident path is intentionally local-shard only. With tensor + parallelism, the regular model loader has already produced each rank's TP + shard, so no DP AllGather is required or desirable here. + """ + + def __init__( + self, + blocks: list[nn.Module], + device: torch.device, + copy_stream: Any, + pin_memory: bool, + ) -> None: + self.device = device + self.copy_stream = copy_stream + self.loaded = False + self._states: list[dict[str, Any]] = [] + self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + + for block in blocks: + params = dict(block.named_parameters()) + bufs = dict(block.named_buffers()) + targets: dict[str, torch.Tensor] = {**params, **bufs} + cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( + params, + bufs, + dp_size=1, + rank=0, + pin_memory=pin_memory, + ) + self._states.append( + { + "targets": targets, + "cpu_shards": cpu_shards, + "metadata": metadata, + } + ) + + def load(self) -> None: + if self.loaded: + return + + # Allocate on the compute stream so the caching allocator can reuse + # blocks released by the encoder stage. Allocating on the copy stream + # would create a separate stream-local pool and inflate peak HBM. + gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + for state in self._states: + block_buffers: dict[torch.dtype, torch.Tensor] = {} + for dtype, cpu_shard in state["cpu_shards"].items(): + block_buffers[dtype] = torch.empty( + cpu_shard.shape, + dtype=dtype, + device=self.device, + ) + gpu_buffers.append(block_buffers) + + self.copy_stream.wait_stream(current_omni_platform.current_stream()) + ready = current_omni_platform.Event() + with current_omni_platform.stream(self.copy_stream): + for state, block_buffers in zip(self._states, gpu_buffers): + for dtype, cpu_shard in state["cpu_shards"].items(): + block_buffers[dtype].copy_(cpu_shard, non_blocking=True) + ready.record(self.copy_stream) + + for state, block_buffers in zip(self._states, gpu_buffers): + targets = state["targets"] + for dtype, metas in state["metadata"].items(): + gpu_buffer = block_buffers[dtype] + for meta in metas: + set_tensor_storage( + targets[meta["name"]], + gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]].view(meta["shape"]), + ) + + current_omni_platform.current_stream().wait_event(ready) + self._gpu_buffers = gpu_buffers + self.loaded = True + + def offload(self) -> None: + if not self.loaded: + return + + # Denoise kernels consume these buffers on the current compute stream. + # Synchronize once at the stage boundary; no D2H copy is necessary + # because the pinned CPU master weights were never overwritten. + current_omni_platform.synchronize() + for state in self._states: + for target in state["targets"].values(): + set_tensor_storage(target, make_offload_placeholder(target)) + self._gpu_buffers.clear() + self.loaded = False + + # ---------------------------------------------------------------------- # # Backend # # ---------------------------------------------------------------------- # @@ -494,6 +616,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): overlapped with computation. """ + _MMAP_PARAM_ATTRS = ("weight_loader", "mmap_weight_transform") + def __init__(self, config: OffloadConfig, device: torch.device): super().__init__(config, device) @@ -504,6 +628,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.rank = 0 self._blocks: list[list[nn.Module]] = [] self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] + self._resident_blocks: list[nn.Module] = [] + self._resident_layer_group: PinnedResidentLayerGroup | None = None + + def load_resident_layers(self) -> None: + """Load the model-declared leading blocks for the denoise stage.""" + if self._resident_layer_group is not None: + self._resident_layer_group.load() + + def offload_resident_layers(self) -> None: + """Release leading blocks before VAE decode to bound peak HBM.""" + if self._resident_layer_group is not None: + self._resident_layer_group.offload() + + def _remember_mmap_param_attrs(self, pipeline: nn.Module) -> None: + """Save loader metadata before ``to_empty`` replaces Parameters.""" + self._mmap_param_attrs = { + name: {attr: value for attr in self._MMAP_PARAM_ATTRS if (value := getattr(param, attr, None)) is not None} + for name, param in pipeline.named_parameters() + } + + def _attach_mmap_param_attrs( + self, + name: str, + replacement: nn.Parameter, + source: nn.Parameter | None = None, + ) -> None: + attrs = getattr(self, "_mmap_param_attrs", {}).get(name, {}) + for attr in self._MMAP_PARAM_ATTRS: + value = attrs.get(attr) + if value is None and source is not None: + value = getattr(source, attr, None) + if value is not None: + setattr(replacement, attr, value) + if attr == "mmap_weight_transform": + replacement.mmap_weight_transform_pending = True def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: """Load DiT weights from safetensors via mmap views (no RSS). @@ -606,6 +765,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): "instead of TP, or disable distributed layerwise offload." ) + # ``Module.to_empty`` creates new Parameter objects and drops custom + # attributes installed by vLLM weight loaders. Preserve them by full + # pipeline name so mmap replacements can reproduce checkpoint layout + # transforms such as MiniMax-H3 grouped QKV. + self._remember_mmap_param_attrs(pipeline) + # # Important: when DiT modules are nested (e.g. transformer contains # transformer.language_model), to_empty("meta") on the parent @@ -730,7 +895,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + self._attach_mmap_param_attrs(name, replacement, param) + parent._parameters[parts[-1]] = replacement loaded += 1 loaded_names.add(name) @@ -785,7 +952,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - parent._parameters[bparts[-1]] = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) + replacement = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) + self._attach_mmap_param_attrs(full_param_name, replacement, bparam) + parent._parameters[bparts[-1]] = replacement block_loaded += 1 loaded_names.add(full_param_name) @@ -896,7 +1065,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + self._attach_mmap_param_attrs(full_name, replacement, param) + parent._parameters[parts[-1]] = replacement loaded += 1 if loaded > 0: @@ -954,24 +1125,74 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): coord.ranks, ) - def _register_on_demand_hook(self, module: nn.Module, label: str) -> None: - """Move *module* to GPU permanently. + def _register_on_demand_hook(self, module: nn.Module, label: str, *, stage_on_demand: bool = False) -> None: + """Prepare a pipeline-managed stage component or keep it resident. - Previously, on-demand hooks moved the module to GPU before forward - and back to CPU after forward (with synchronize + empty_cache). - However, _post_forward's global synchronize() and empty_cache() - disrupt the async prefetch pipeline (copy_stream / comm_stream) - used by distributed layerwise offload, causing garbled output. - - Moving the module to GPU permanently avoids this issue at the - cost of slightly higher HBM usage (VAE/encoders are small). + Components that expose an explicit stage lifecycle are initially + offloaded and loaded by their pipeline only around encode/decode. + Other models retain the conservative resident behavior because a + generic post-forward hook can disrupt the DiT prefetch streams. """ + offload_to_cpu = getattr(module, "offload_to_cpu", None) + if stage_on_demand and callable(offload_to_cpu): + offload_to_cpu() + logger.info("Prepared %s (%s) for pipeline-managed staged offload", label, module.__class__.__name__) + return module.to(self.device) logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__) + def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool: + """Stream plan-declared encoder blocks on each rank without AllGather.""" + if plan is None or name not in plan.encoder_block_attrs: + return False + if getattr(module, "_omni_layerwise_enabled", False): + return True + + from operator import attrgetter + + from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook + + hooks = [] + block_groups = [] + copy_stream = current_omni_platform.Stream() + for block_path in plan.encoder_block_attrs[name]: + try: + blocks = attrgetter(block_path)(module) + except AttributeError: + logger.warning("Encoder offload path %s.%s was not found", name, block_path) + continue + if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1: + logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path) + continue + group_hooks = [ + apply_block_hook(blocks[-1], blocks[0], self.device, copy_stream, self.config.pin_cpu_memory) + ] + group_hooks.extend( + apply_block_hook(block, blocks[index + 1], self.device, copy_stream, self.config.pin_cpu_memory) + for index, block in enumerate(blocks[:-1]) + ) + for index, hook in enumerate(group_hooks): + hook._prev_hook = group_hooks[index - 1] + hooks.extend(group_hooks) + block_groups.append(blocks) + + if not hooks: + return False + # The component lifecycle uses these generic attributes to keep only + # non-block encoder state resident during the encode phase. + module._omni_layerwise_hooks = hooks + module._omni_layerwise_block_groups = block_groups + module._omni_layerwise_enabled = True + logger.info( + "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", + name, + sum(len(blocks) for blocks in block_groups), + len(block_groups), + ) + return True + def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool: """Try to apply layerwise offload to a large submodule's blocks. - Resolution order: 1. OffloadPlan.offload_submodules (declarative, if plan is provided) 2. Heuristic search for common block-list attributes @@ -1081,6 +1302,78 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._blocks.append(blocks) return True + def _prepare_dit_non_block_modules( + self, + dit_module: nn.Module, + dit_name: str, + blocks_attr_names: list[str], + all_dit_modules: set[int], + plan: OffloadPlan | None, + ) -> None: + """Place or hook the DiT parts that are outside its repeated blocks. + + This must run even when every repeated block is resident. Otherwise + an all-resident stage skips placement for modules such as H3's token + refiner and enters the forward pass with CPU or meta tensors. + """ + _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB + for name, module in dit_module.named_children(): + if name in blocks_attr_names: + logger.debug("Skipped blocks module %s", name) + continue + + has_meta = any(getattr(param, "is_meta", False) for param in module.parameters()) + if has_meta: + self._load_module_weights_from_mmap(module, dit_name, name) + module_mb = ( + sum( + param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0 + for param in module.parameters() + ) + / 1048576 + ) + explicitly_planned = plan is not None and name in plan.offload_submodules + if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD: + if id(module) in all_dit_modules: + logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) + elif self._try_layerwise_offload_submodule(module, name, plan): + pass + else: + self._register_on_demand_hook(module, name) + continue + + try: + module.to(self.device) + except (NotImplementedError, RuntimeError): + self._load_module_weights_from_mmap(module, dit_name, name) + # Non-persistent buffers such as RoPE frequencies do not + # exist in the checkpoint and must be reconstructed. + has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True)) + if has_meta_buffer: + saved_params = { + param_name: param.data.clone() + for param_name, param in module.named_parameters() + if not getattr(param, "is_meta", False) + } + module.to_empty(device=self.device) + for submodule in module.modules(): + if hasattr(submodule, "reset_parameters"): + submodule.reset_parameters() + for param_name, param in module.named_parameters(): + if param_name in saved_params: + param.data.copy_(saved_params[param_name]) + try: + module.to(self.device) + except Exception: + logger.warning("Module %s still has meta params after mmap load", name) + + for param in dit_module._parameters.values(): + if param is not None and not getattr(param, "is_meta", False): + param.data = param.data.to(self.device, non_blocking=True) + for buffer in dit_module._buffers.values(): + if buffer is not None: + buffer.data = buffer.data.to(self.device, non_blocking=True) + def enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("DistributedLayerwiseOffloadBackend already enabled") @@ -1102,7 +1395,19 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # When present, replaces heuristic block discovery. plan = get_offload_plan(pipeline) + if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths): + logger.warning( + "dlo_resident_layers=%d was requested, but this model declares no " + "resident_dit_paths; all blocks will be streamed.", + self.config.dlo_resident_layers, + ) + # Load weights via mmap for DLO+AllGather. + # TODO(offload): Add a rank-local mmap path for dlo_no_use_allgather. + # It must apply declarative checkpoint-to-runtime adapters before the + # standard TP loader shards each block, and must cover staged encoders + # and VAEs as well as DiT blocks. The current AllGather mmap path is + # deliberately not reused because it changes the weight layout. # Gate condition MUST match diffusers_loader.py: # supports_mmap_loading(pipeline) and not _has_online_quant # When the gate is False, the loader has already loaded weights via @@ -1129,10 +1434,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) # during the DiT forward pass. They are only needed briefly for # text-encoding (before DiT) and VAE-decode (after DiT). - for enc in modules.encoders: - self._register_on_demand_hook(enc, "encoder") - for vae in modules.vaes: - self._register_on_demand_hook(vae, "vae") + for enc, enc_name in zip(modules.encoders, modules.encoder_names): + self._try_layerwise_offload_encoder(enc, enc_name, plan) + self._register_on_demand_hook( + enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths + ) + for vae, vae_name in zip(modules.vaes, modules.vae_names): + self._register_on_demand_hook( + vae, + "vae", + stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths), + ) # Move resident modules to GPU (small modules needed every forward) for name, module in zip(modules.resident_names, modules.resident_modules): @@ -1163,93 +1475,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): dit_module.to(self.device) continue + resident_count = 0 + if plan is not None and dit_name in plan.resident_dit_paths: + resident_count = min(self.config.dlo_resident_layers, len(blocks)) + if resident_count: + resident_blocks = blocks[:resident_count] + self._resident_blocks.extend(resident_blocks) + blocks = blocks[resident_count:] + logger.info( + "Keeping %d leading blocks resident on %s; streaming %d tail blocks", + resident_count, + dit_name, + len(blocks), + ) + + self._prepare_dit_non_block_modules( + dit_module, + dit_name, + blocks_attr_names, + all_dit_modules, + plan, + ) + num_blocks = len(blocks) + if num_blocks == 0: + logger.info("All blocks for %s are resident; no streaming hooks required", dit_name) + continue if num_blocks <= 1: logger.warning( - "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", dit_name, dit_module.__class__.__name__, ) - dit_module.to(self.device) + self._resident_blocks.extend(blocks) continue - # Move non-block modules to GPU (they stay resident). - # Large modules (> 1 GB) use on-demand hooks instead — they are - # only needed briefly (e.g. language_model for text encoding) - # and keeping them resident wastes HBM during the DiT loop. - _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB # alias for readability - for name, m in dit_module.named_children(): - if name not in blocks_attr_names: - _has_meta = any((p.is_meta if hasattr(p, "is_meta") else False) for p in m.parameters()) - if _has_meta: - # Meta params: load from safetensors before moving to GPU - self._load_module_weights_from_mmap(m, dit_name, name) - _mb = ( - sum( - (p.nelement() * p.element_size() if not (hasattr(p, "is_meta") and p.is_meta) else 0) - for p in m.parameters() - ) - / 1048576 - ) - if _mb > _ON_DEMAND_THRESHOLD: - # Skip if this submodule is already a separate DiT - # module (e.g. transformer.language_model when both - # "transformer.language_model" and "transformer" are - # in _dit_modules) — hooks already applied. - if id(m) in all_dit_modules: - logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) - elif self._try_layerwise_offload_submodule(m, name, plan): - pass # layerwise hooks applied - else: - self._register_on_demand_hook(m, name) - else: - try: - m.to(self.device) - except (NotImplementedError, RuntimeError): - # Meta params: load from safetensors via mmap, then move to GPU - self._load_module_weights_from_mmap(m, dit_name, name) - # Materialize any remaining meta buffers (e.g. RoPE - # inv_freq, timestep freqs) on device and - # recompute them from their original formulas. - # These are non-persistent buffers not present in - # the checkpoint — they must be reconstructed, not - # left as torch.empty. - # NOTE: Do NOT call reset_parameters() here — it - # would re-initialize ALL parameters (including - # mmap-loaded weights) with random values. - _has_meta_buf = any(hasattr(b, "is_meta") and b.is_meta for b in m.buffers(recurse=True)) - if _has_meta_buf: - # Save mmap-loaded param values before - # reset_parameters (which re-initializes - # ALL params with random values). - saved_params = { - n: p.data.clone() - for n, p in m.named_parameters() - if not (hasattr(p, "is_meta") and p.is_meta) - } - m.to_empty(device=self.device) - for submod in m.modules(): - if hasattr(submod, "reset_parameters"): - submod.reset_parameters() - # Restore mmap-loaded weights - for n, p in m.named_parameters(): - if n in saved_params: - p.data.copy_(saved_params[n]) - try: - m.to(self.device) - except Exception: - logger.warning("Module %s still has meta params after mmap load", name) - else: - logger.debug(f"Skipped blocks module {name}") - - # Move top-level params/buffers to GPU - for param in dit_module._parameters.values(): - if param is not None and not (hasattr(param, "is_meta") and param.is_meta): - param.data = param.data.to(self.device, non_blocking=True) - for buffer in dit_module._buffers.values(): - if buffer is not None: - buffer.data = buffer.data.to(self.device, non_blocking=True) - # Register hooks in a circular sliding window: # last block prefetches first block, block i prefetches block (i+1) # All hooks share 2 global device buffers (RFC: "exactly two layers on device") @@ -1306,7 +1566,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups.append(block_hooks) self._blocks.append(blocks) + if self._resident_blocks: + self._resident_layer_group = PinnedResidentLayerGroup( + self._resident_blocks, + self.device, + self.copy_stream, + self.config.pin_cpu_memory, + ) + pipeline._dlo_residency_controller = self + if not self._all_hook_groups: + self.enabled = bool(self._resident_blocks) return # Unified allocation: 2 shared output buffers + 2 shared shard buffers @@ -1416,8 +1686,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._on_demand_handles = [] self._on_demand_shard_infos = [] + self.offload_resident_layers() self._blocks.clear() self._all_hook_groups.clear() + self._resident_blocks.clear() + self._resident_layer_group = None self.enabled = False logger.info("Distributed layer-wise offloading disabled") diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index d6fcb2569..b93c9ed73 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -19,6 +19,7 @@ class PipelineModules: encoders: list[nn.Module] encoder_names: list[str] vaes: list[nn.Module] + vae_names: list[str] = field(default_factory=list) resident_modules: list[nn.Module] = field(default_factory=list) resident_names: list[str] = field(default_factory=list) @@ -153,7 +154,7 @@ class ModuleDiscovery: dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared) encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared) - vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) + vaes, vae_names = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared) return PipelineModules( @@ -163,5 +164,6 @@ class ModuleDiscovery: encoder_names=encoder_names, vaes=vaes, resident_modules=residents, + vae_names=vae_names, resident_names=resident_names, ) diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py new file mode 100644 index 000000000..da32ab520 --- /dev/null +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""On-demand module staging backed by an immutable pinned CPU snapshot.""" + +from __future__ import annotations + +from itertools import chain +from typing import Any + +import torch +from torch import nn + +from vllm_omni.platforms import current_omni_platform + +from .tensor_utils import set_tensor_storage + + +class PinnedModuleStager: + """Stage an immutable module without copying device weights back to CPU. + + ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter + after each forward. In inference the weights are immutable, so retain one + pinned CPU master instead. ``load`` materializes fresh device tensors from + that master; ``offload`` only rebinds the Parameters/Buffers to it and + releases the device allocations. + """ + + def __init__( + self, + module: nn.Module, + device: torch.device, + *, + pin_memory: bool = True, + copy_stream: Any | None = None, + ) -> None: + self.device = device + self.copy_stream = copy_stream or current_omni_platform.Stream() + self.loaded = False + self._entries: list[tuple[torch.Tensor, torch.Tensor]] = [] + self._device_tensors: list[torch.Tensor] = [] + + # parameters()/buffers() remove duplicate objects by default, which + # preserves tied weights and buffer aliases when their storage moves. + for target in chain(module.parameters(), module.buffers()): + master = target.detach() + if master.device.type != "cpu": + master = master.to("cpu") + if pin_memory and not master.is_pinned(): + master = master.pin_memory() + set_tensor_storage(target, master) + self._entries.append((target, master)) + + def load(self) -> None: + if self.loaded: + return + + # Allocate on the compute stream so the caching allocator can reuse + # memory released by the preceding encoder/DiT stage. + device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries] + compute_stream = current_omni_platform.current_stream() + self.copy_stream.wait_stream(compute_stream) + ready = current_omni_platform.Event() + with current_omni_platform.stream(self.copy_stream): + for device_tensor, (_, master) in zip(device_tensors, self._entries): + device_tensor.copy_(master, non_blocking=master.is_pinned()) + ready.record(self.copy_stream) + + for device_tensor, (target, _) in zip(device_tensors, self._entries): + set_tensor_storage(target, device_tensor) + compute_stream.wait_event(ready) + self._device_tensors = device_tensors + self.loaded = True + + def offload(self) -> None: + if not self.loaded: + return + + # The module has completed on the compute stream. Synchronize once at + # the stage boundary, then release HBM by rebinding to the unchanged + # host masters. No device-to-host transfer is necessary. + current_omni_platform.synchronize() + for target, master in self._entries: + set_tensor_storage(target, master) + self._device_tensors.clear() + self.loaded = False + current_omni_platform.empty_cache() + + +__all__ = ["PinnedModuleStager"] diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 575194ea0..f8943757e 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -35,10 +35,21 @@ class OffloadPlan: for large non-DiT submodules within a DiT that should be independently offloaded with their own hooks. e.g. ``{"context_encoder": "layers"}`` + resident_dit_paths: DiT paths whose leading blocks may be kept on the + device when ``dlo_resident_layers`` is nonzero. Keeping this + model-declared avoids applying a consumer-GPU tuning knob to + auxiliary or dual DiTs unintentionally. + encoder_block_attrs: Maps encoder paths to rank-local block-list paths. + These blocks are streamed with ordinary layerwise hooks, never + with the DiT AllGather group. """ + on_demand_component_paths: frozenset[str] = field(default_factory=frozenset) + block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) offload_submodules: dict[str, str] = field(default_factory=dict) + resident_dit_paths: frozenset[str] = field(default_factory=frozenset) + encoder_block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: @@ -58,4 +69,6 @@ def supports_mmap_loading(pipeline: nn.Module) -> bool: (to decide whether to call ``_load_weights_via_mmap``), ensuring both paths use the same gate condition. """ - return any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()) + return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any( + callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules() + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 6d8f40e0f..a45f3bdee 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -91,6 +91,11 @@ _DIFFUSION_MODELS = { "pipeline_minimax_h3", "MiniMaxH3Pipeline", ), + "MiniMaxH3ModularPipeline": ( + "minimax_h3", + "pipeline_minimax_h3", + "MiniMaxH3Pipeline", + ), "StableAudioPipeline": ( "stable_audio", "pipeline_stable_audio", @@ -427,9 +432,11 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - if sp_size <= 1: return - # Find transformer model(s) in the pipeline that have _sp_plan - # Include transformer_2 for two-stage models (e.g., Wan MoE) - transformer_attrs = ["transformer", "transformer_2", "dit", "unet"] + # Prefer the pipeline's declared DiT components so custom component + # names receive the same SP hooks as conventional transformer names. + transformer_attrs = getattr(model, "_dit_modules", None) + if not transformer_attrs: + transformer_attrs = ("transformer", "transformer_2", "dit", "unet") applied_count = 0 for attr in transformer_attrs: @@ -508,6 +515,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func", "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", + "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", "SoulXSingerPipeline": "get_soulxsinger_post_process_func", "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func", diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index dad02583f..78d8c6536 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/vllm_omni/diffusion/utils/hf_utils.py @@ -7,6 +7,24 @@ from vllm.transformers_utils.config import get_hf_file_to_dict logger = init_logger(__name__) +DIFFUSION_MODEL_INDEX_FILES = ( + "model_index.json", + "modular_model_index.json", +) + + +def get_diffusion_model_index( + model_name: str, + *, + revision: str | None = None, +) -> dict | None: + """Read the first standard Diffusers pipeline index available.""" + for filename in DIFFUSION_MODEL_INDEX_FILES: + config = get_hf_file_to_dict(filename, model_name, revision=revision) + if isinstance(config, Mapping): + return dict(config) + return None + def load_diffusers_config(model_name) -> dict: from diffusers.pipelines.pipeline_utils import DiffusionPipeline @@ -58,33 +76,32 @@ def is_diffusion_model(model_name: str) -> bool: """Check if a model is a diffusion model. Uses multiple fallback strategies to detect diffusion models: - 1. Check local file system for model_index.json (fastest, no imports) + 1. Check the local file system for a standard Diffusers index 2. Check using vllm's get_hf_file_to_dict utility 3. Try the standard diffusers approach (may fail due to import issues) """ # Strategy 1: Check local file system first (fastest, avoids import issues) if os.path.isdir(model_name): - model_index_path = os.path.join(model_name, "model_index.json") - if os.path.exists(model_index_path): + for filename in DIFFUSION_MODEL_INDEX_FILES: + model_index_path = os.path.join(model_name, filename) + if not os.path.exists(model_index_path): + continue try: import json with open(model_index_path) as f: config_dict = json.load(f) if config_dict.get("_class_name") and config_dict.get("_diffusers_version"): - logger.debug("Detected diffusion model via local model_index.json") + logger.debug("Detected diffusion model via local %s", filename) return True except Exception as e: - logger.debug("Failed to read local model_index.json: %s", e) + logger.debug("Failed to read local %s: %s", filename, e) # Strategy 2: Check using vllm's utility (works for both local and remote models) - try: - config_dict = get_hf_file_to_dict("model_index.json", model_name) - if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"): - logger.debug("Detected diffusion model via model_index.json") - return True - except Exception as e: - logger.debug("Failed to check model_index.json via get_hf_file_to_dict: %s", e) + config_dict = get_diffusion_model_index(model_name) + if config_dict is not None and config_dict.get("_class_name") and config_dict.get("_diffusers_version"): + logger.debug("Detected diffusion model via a standard Diffusers index") + return True # Strategy 3: Try the standard diffusers approach (may fail due to import issues) # This is last because it requires importing diffusers/xformers/flash_attn diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 3255e59db..2e8111a31 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -280,8 +280,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): exc, ) else: - self._compile_transformer("transformer") - self._compile_transformer("transformer_2") + transformer_attrs = getattr(self.pipeline, "_dit_modules", None) + if not transformer_attrs: + transformer_attrs = ("transformer", "transformer_2") + for attr_name in transformer_attrs: + self._compile_transformer(attr_name) else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fbc2bec3a..633eec436 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -140,8 +140,9 @@ class OmniEngineArgs(EngineArgs): stage_connector_spec: Extra configuration for stage connector async_chunk: If set to True, perform async chunk worker_type: Model Type, e.g., "ar" or "generation" - task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base). - If not specified, will be inferred from model path. + task_type: Model-defined startup task type. Consumers validate the + supported values and decide whether it selects request behavior, + task-specific weights, or both. omni_master_address: TCP address that the OmniMasterServer (running inside AsyncOmniEngine) listens on for engine core registrations. Required when single-stage mode is active. @@ -495,6 +496,7 @@ class OrchestratorArgs: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_resident_layers: int = 0 boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index b3307c3f5..c938d0757 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1019,6 +1019,7 @@ class AsyncOmniEngine: "max_num_seqs": kwargs.get("max_num_seqs") or 1, "parallel_config": parallel_config, "model_class_name": kwargs.get("model_class_name", None), + "task_type": kwargs.get("task_type", None), "model_config": kwargs.get("model_config", None), "additional_config": kwargs.get("additional_config", None), "step_execution": kwargs.get("step_execution", False), @@ -1032,6 +1033,7 @@ class AsyncOmniEngine: "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), + "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 3638b791e..e8284ba68 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -8,6 +8,7 @@ out of StageEngineCoreClient into reusable functions. from __future__ import annotations +import copy import fcntl import importlib import multiprocessing as mp @@ -15,7 +16,7 @@ import os import time from collections.abc import Callable, Generator, Sequence from contextlib import contextmanager -from dataclasses import dataclass, replace +from dataclasses import dataclass, fields, replace from typing import Any, Literal, cast from vllm.logger import init_logger @@ -25,7 +26,10 @@ from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.executor import Executor -from vllm_omni.config.omni_config import BaseVllmOmniStageConfig +from vllm_omni.config.omni_config import ( + BaseVllmOmniStageConfig, + VllmOmniDiffusionStageConfig, +) from vllm_omni.config.stage_config import StageType from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs @@ -348,6 +352,30 @@ class StageMetadata: replica_id: int = 0 +def _apply_rocm_attention_backend( + engine_args: dict[str, Any], + stage_type: str | StageType, +) -> None: + """Preserve Omni's ROCm attention-backend compatibility default.""" + if ( + not current_omni_platform.is_rocm() + or stage_type == StageType.DIFFUSION + or engine_args.get("attention_backend") is not None + ): + return + + from vllm._aiter_ops import rocm_aiter_ops + + if rocm_aiter_ops.is_enabled(): + engine_args["attention_backend"] = "ROCM_AITER_FA" + # Before vLLM v0.19.0, the default attention backend is TRITON_ATTN for ROCm. + # Since vLLM v0.19.0, the default attention backend is ROCM_ATTN for ROCm. + # However, the compatibility of ROCM_ATTN with Omni is not guaranteed. + # Therefore, we still use TRITON_ATTN as the default attention backend, + # when the selected_backend is not specified. + engine_args["attention_backend"] = "TRITON_ATTN" + + def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: """Extract metadata through the active production legacy path. @@ -358,18 +386,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: stage_type: Literal["llm", "diffusion"] = _get_attr_or_item(stage_config, "stage_type", "llm") engine_args = stage_config.engine_args - if current_omni_platform.is_rocm(): - if stage_type != "diffusion" and engine_args.get("attention_backend") is None: - from vllm._aiter_ops import rocm_aiter_ops - - if rocm_aiter_ops.is_enabled(): - engine_args["attention_backend"] = "ROCM_AITER_FA" - # Before vLLM v0.19.0, the default attention backend is TRITON_ATTN for ROCm. - # Since vLLM v0.19.0, the default attention backend is ROCM_ATTN for ROCm. - # However, the compatibility of ROCM_ATTN with Omni is not guaranteed. - # Therefore, we still use TRITON_ATTN as the default attention backend, - # when the selected_backend is not specified. - engine_args["attention_backend"] = "TRITON_ATTN" + _apply_rocm_attention_backend(engine_args, stage_type) runtime_cfg = stage_config.runtime engine_input_source: list[int] = _get_attr_or_item(stage_config, "engine_input_source", []) @@ -748,25 +765,110 @@ def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None, os.environ[key] = old_value -def build_engine_args_dict( - stage_config: Any, - model: str, - stage_connector_spec: dict[str, Any] | None = None, - cli_tokenizer: str | None = None, +def _project_omni_config_fields( + config: Any, + *, + exclude: frozenset[str] = frozenset(), ) -> dict[str, Any]: - """Build the normalized engine args dict for one stage.""" - engine_args = stage_config.engine_args - # HACK (Alex) Tensor parallel size should not be passed as None; - # remove it if this is the case so that we fall back to default - # creation from vLLM's engine args. - # NOTE: This will be fixed more generically in ongoing work for engine arg filtering. - if "tensor_parallel_size" in engine_args and engine_args["tensor_parallel_size"] is None: - del engine_args["tensor_parallel_size"] + """Copy defined typed config fields into backend adapter kwargs.""" + projected: dict[str, Any] = {} + for config_field in fields(config): + name = config_field.name + if name in exclude: + continue + value = getattr(config, name) + if value is not None: + projected[name] = copy.deepcopy(value) + return projected - stage_type = _get_attr_or_item(stage_config, "stage_type", "llm") - stage_id = stage_config.stage_id - engine_args_dict = _to_dict(engine_args) +def _project_omni_stage_engine_args( + stage_config: BaseVllmOmniStageConfig, +) -> dict[str, Any]: + """Read backend inputs from one structured stage config.""" + engine_args: dict[str, Any] = {} + is_diffusion = isinstance(stage_config, VllmOmniDiffusionStageConfig) + + if is_diffusion: + engine_args.update(_project_omni_config_fields(stage_config.diffusion_config)) + + for config, excluded_fields in ( + ( + stage_config.model_config, + frozenset({"default_sampling_params", "has_sampling_extra_args"}), + ), + (stage_config.load_config, frozenset()), + (stage_config.cache_config, frozenset()), + (stage_config.scheduler_config, frozenset()), + ( + stage_config.runtime_config, + frozenset({"devices", "num_replicas", "env", "num_gpus"}), + ), + ): + engine_args.update( + _project_omni_config_fields( + config, + exclude=excluded_fields, + ) + ) + + # The legacy builder always emits this key, including for pipelines such + # as Audex that intentionally defer architecture discovery to HF config. + engine_args["model_arch"] = copy.deepcopy(stage_config.model_config.model_arch) + + topology = stage_config.stage_pipeline_config + topology_engine_args = { + "model_stage": stage_config.model_stage, + "worker_type": stage_config.worker_type, + "scheduler_cls": stage_config.scheduler_cls, + "hf_config_name": stage_config.hf_config_name, + "engine_output_type": stage_config.engine_output_type, + "custom_process_next_stage_input_func": stage_config.custom_process_next_stage_input_func, + "retains_state_across_chunks": topology.retains_state_across_chunks, + } + engine_args.update( + {name: copy.deepcopy(value) for name, value in topology_engine_args.items() if value is not None} + ) + + connector_config = stage_config.connector_config + engine_args["async_chunk"] = connector_config.async_chunk + if connector_config.omni_kv_config is not None: + engine_args["omni_kv_config"] = copy.deepcopy(connector_config.omni_kv_config) + + if is_diffusion: + engine_args["parallel_config"] = _project_omni_config_fields( + stage_config.parallel_config, + exclude=frozenset({"world_size"}), + ) + else: + engine_args.update( + _project_omni_config_fields( + stage_config.parallel_config, + exclude=frozenset({"world_size"}), + ) + ) + + quantization_config = stage_config.quantization_config + if quantization_config is not None: + quantization_key = ( + "quantization" if isinstance(quantization_config, str) and not is_diffusion else "quantization_config" + ) + engine_args[quantization_key] = copy.deepcopy(quantization_config) + + return engine_args + + +def _finalize_engine_args_dict( + engine_args_dict: dict[str, Any], + *, + stage_type: str | StageType, + stage_id: int, + model: str, + stage_connector_spec: dict[str, Any] | None, + cli_tokenizer: str | None, + has_sampling_extra_args: bool, +) -> dict[str, Any]: + """Apply representation-independent engine adapter behavior.""" pipeline_model_root = model model = engine_args_dict.pop("model", None) or model stage_defines_tokenizer = ( @@ -826,15 +928,15 @@ def build_engine_args_dict( if stage_connector_spec: engine_args_dict["stage_connector_spec"] = dict(stage_connector_spec or {}) - if stage_type == "diffusion": + is_diffusion = stage_type == StageType.DIFFUSION + if is_diffusion: from vllm_omni.diffusion.data import parse_attention_config if engine_args_dict.get("diffusion_attention_config") is not None: engine_args_dict["diffusion_attention_config"] = parse_attention_config( - engine_args_dict.get("diffusion_attention_config"), + engine_args_dict["diffusion_attention_config"], ) - - if stage_type != "diffusion": + else: resolve_worker_cls(engine_args_dict) if engine_args_dict.get("worker_type") == "generation": @@ -844,17 +946,86 @@ def build_engine_args_dict( engine_args_dict.setdefault("disable_hybrid_kv_cache_manager", True) engine_args_dict.setdefault("enable_prefix_caching", False) - # Check whether the stage's default_sampling_params defines extra_args. - default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {})) - engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args")) + engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args # TODO: Remove this after the performance regression is fixed # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression _maybe_set_qwen3_omni_moe_env(engine_args_dict) - return engine_args_dict +def build_legacy_engine_args_dict( + stage_config: Any, + model: str, + stage_connector_spec: dict[str, Any] | None = None, + cli_tokenizer: str | None = None, +) -> dict[str, Any]: + """Implement engine-argument building for the legacy stage representation.""" + engine_args_dict = _to_dict(stage_config.engine_args) + # Legacy configs can materialize an omitted optional TP size as None. + # Remove it from the detached adapter dict so the backend default applies + # without mutating stage_config.engine_args. + if engine_args_dict.get("tensor_parallel_size") is None: + engine_args_dict.pop("tensor_parallel_size", None) + + default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {})) + return _finalize_engine_args_dict( + engine_args_dict, + stage_type=_get_attr_or_item(stage_config, "stage_type", "llm"), + stage_id=stage_config.stage_id, + model=model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=cli_tokenizer, + has_sampling_extra_args=bool(default_sp.get("extra_args")), + ) + + +def build_engine_args_dict( + stage_config: Any, + model: str, + stage_connector_spec: dict[str, Any] | None = None, + cli_tokenizer: str | None = None, +) -> dict[str, Any]: + """Build engine arguments through the stable production entry point. + + Production inputs still use the legacy stage representation. Keep that + compatibility choice behind this function so callers do not bind directly + to a representation-specific implementation. + """ + return build_legacy_engine_args_dict( + stage_config, + model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=cli_tokenizer, + ) + + +def build_engine_args_dict_from_omni_stage_config( + stage_config: BaseVllmOmniStageConfig, + model: str, + stage_connector_spec: dict[str, Any] | None = None, + cli_tokenizer: str | None = None, +) -> dict[str, Any]: + """Project one typed stage config into backend engine arguments. + + This projection is prepared for the RFC #4021 stage-init cutover. Current + production startup reaches the legacy implementation through + ``build_engine_args_dict`` while strategy and startup-plan inputs still + use the legacy representation. + """ + engine_args_dict = _project_omni_stage_engine_args(stage_config) + _apply_rocm_attention_backend(engine_args_dict, stage_config.stage_type) + return _finalize_engine_args_dict( + engine_args_dict, + stage_type=stage_config.stage_type, + stage_id=stage_config.stage_id, + model=model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=cli_tokenizer, + has_sampling_extra_args=stage_config.model_config.has_sampling_extra_args, + ) + + def build_vllm_config( stage_config: Any, model: str, diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 8f13a3643..bf84f7793 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -370,6 +370,9 @@ class StageRuntime: executor_class = None engine_args_dict = None if base_metadata.stage_type != "diffusion": + # The stable adapter entry point still receives the same + # legacy stage object as replica planning. Its implementation + # switches only at the coordinated RFC #4021 cutover. engine_args_dict = build_engine_args_dict( stage_cfg, self._model, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index df3d8aee4..8256d084a 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -213,9 +213,10 @@ class OmniServeCommand(CLISubcommand): "--task-type", type=str, default=None, - choices=["CustomVoice", "VoiceDesign", "Base"], - help="Default task type for TTS models (CustomVoice, VoiceDesign, or Base). " - "If not specified, will be inferred from model path.", + help="Model-defined startup task type. The selected model validates " + "supported values; for example, TTS models accept CustomVoice, " + "VoiceDesign, or Base, while diffusion models may use it to select " + "task-specific weights. If omitted, the model default is used.", ) # Forced aligner / word timestamps. --forced-aligner is the opt-in # toggle; heavier knobs (gpu_memory_utilization, dtype, max_model_len) @@ -646,15 +647,25 @@ class OmniServeCommand(CLISubcommand): action="store_true", default=True, help="Use shard + AllGather for weight reconstruction (default: True). " - "When disabled (--dlo-no-use-allgather), each rank loads full weights " - "via H2D only — no sharding, no AllGather, no concurrent request " - "requirement, but N× CPU memory.", + "When disabled (--dlo-no-use-allgather), each rank streams the " + "standard loader's rank-local tensors via H2D only — no additional " + "DP sharding, no AllGather, and no concurrent-request requirement.", ) omni_config_group.add_argument( "--dlo-no-use-allgather", dest="dlo_use_allgather", action="store_false", - help="Disable AllGather: each rank loads full weights independently.", + help=( + "Disable AllGather and stream standard-loader rank-local weights " + "independently (including existing TP shards)." + ), + ) + omni_config_group.add_argument( + "--dlo-resident-layers", + type=int, + default=0, + help="Keep this many leading main-DiT blocks resident on the device " + "while distributed layerwise offload streams the remaining blocks.", ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( @@ -924,6 +935,9 @@ def run_headless(args: TrackingNamespace) -> None: # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally # bracket each replica's spawn below with setup_stage_devices so they # don't all stack on cuda:0 (see ``per_replica_devices`` above). + # Headless startup still supplies the legacy OmegaConf stage shape through + # the stable adapter entry point. The implementation switches only when + # RFC #4021 threads structured stage configs through the launch plan. engine_args_dict = build_engine_args_dict( stage_cfg, model, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 55ebffa02..3cae33661 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -9,6 +9,7 @@ from typing import Any, Literal import huggingface_hub import vllm.envs as envs from vllm.logger import init_logger +from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -73,6 +74,16 @@ def omni_snapshot_download(model_id: str) -> str: return snapshot_download(model_id) + # Modular Diffusers repositories describe independently loadable + # components. Let the selected pipeline download only its component + # sources instead of eagerly materializing the entire repository. + try: + if file_or_path_exists(model_id, "modular_model_index.json", revision=None): + return model_id + except (huggingface_hub.errors.GatedRepoError, huggingface_hub.errors.RepositoryNotFoundError): + # Preserve the more helpful errors raised by the full-download path. + pass + try: download_weights_from_hf_specific( model_name_or_path=model_id, diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 40c801f14..5a4f9f826 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -481,18 +481,32 @@ def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: if not frames: raise ValueError("No frames found to encode.") - frames_np = np.stack(frames, axis=0) - if frames_np.ndim == 4 and frames_np.shape[-1] == 4: - frames_np = frames_np[..., :3] + frame_shape = frames[0].shape + has_alpha = len(frame_shape) == 3 and frame_shape[-1] == 4 + output_shape = (*frame_shape[:-1], 3) if has_alpha else frame_shape + frames_u8 = np.empty((len(frames), *output_shape), dtype=np.uint8) + common_dtype = np.result_type(*(frame.dtype for frame in frames)) - if frames_np.dtype == np.uint8: - frames_u8 = frames_np - else: - frames_np = np.clip(frames_np, 0.0, 1.0) - frames_np *= 255.0 - frames_u8 = np.round(frames_np).astype(np.uint8) + # Convert one frame at a time instead of stacking the normalized float + # payload first. Long videos can otherwise require another full-size + # float array plus conversion temporaries before encoding. + for index, frame in enumerate(frames): + if frame.shape != frame_shape: + raise ValueError("All video frames must have the same shape.") + frame = frame[..., :3] if has_alpha else frame + if frame.dtype == np.uint8: + frames_u8[index] = frame + continue - return np.ascontiguousarray(frames_u8) + # np.stack(), used by the previous implementation, promoted mixed + # frame dtypes before scaling. Preserve those rounding semantics + # without allocating a full-video float buffer. + scaled = np.clip(frame.astype(common_dtype, copy=False), 0.0, 1.0) + scaled *= 255.0 + np.rint(scaled, out=scaled) + frames_u8[index] = scaled + + return frames_u8 def _encode_video_bytes( diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index dab544d62..8a8f34c15 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -16,7 +16,10 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import _DEPLOY_DIR from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs -from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero +from vllm_omni.diffusion.utils.hf_utils import ( + _looks_like_dreamzero, + get_diffusion_model_index, +) from vllm_omni.entrypoints.stage_utils import _to_dict from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.platforms import current_omni_platform @@ -74,9 +77,9 @@ def _try_get_class_name_from_diffusers_config(model: str) -> str | None: Returns: Model type string if found, None otherwise """ - model_index = get_hf_file_to_dict("model_index.json", model, revision=None) + model_index = get_diffusion_model_index(model) if model_index and isinstance(model_index, dict) and "_class_name" in model_index: - logger.debug(f"Found model_type '{model_index['_class_name']}' in model_index.json") + logger.debug(f"Found Diffusers model type '{model_index['_class_name']}'") return model_index["_class_name"] return None @@ -278,12 +281,12 @@ def resolve_model_config_path(model: str) -> str | None: model_type = hf_config.model_type except (ValueError, Exception): # If standard transformers format fails, try diffusers format - if file_or_path_exists(model, "model_index.json", revision=None): + if get_diffusion_model_index(model) is not None: model_type = _try_get_class_name_from_diffusers_config(model) if model_type is None: raise ValueError( f"Could not determine model_type for diffusers model: {model}. " - f"Please ensure the model has 'model_type' in transformer/config.json or model_index.json" + "Please ensure its Diffusers pipeline index contains '_class_name'" ) elif file_or_path_exists(model, "config.json", revision=None): # Try to read config.json manually for custom models like Bagel that fail get_config @@ -308,7 +311,7 @@ def resolve_model_config_path(model: str) -> str | None: if model_type is None: raise ValueError( f"Could not determine model_type for model: {model}. " - f"Model is not in standard transformers format and does not have model_index.json. " + "Model is not in standard transformers or Diffusers format. " f"Please ensure the model has proper configuration files with 'model_type' field" ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 4671e3b96..a182848cf 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -770,8 +770,11 @@ AutoImageProcessor.register(MiniCPMOConfig, MiniCPMVImageProcessor, exist_ok=Tru # ============== SigLIP Vision Transformer Classes ============== +FlashAttentionUnpaddingMetadata: TypeAlias = tuple[torch.Tensor, torch.Tensor, int] + + # Copied from transformers.models.llama.modeling_llama._get_unpad_data -def _get_unpad_data(attention_mask): +def _get_unpad_data(attention_mask: torch.Tensor) -> FlashAttentionUnpaddingMetadata: seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() max_seqlen_in_batch = seqlens_in_batch.max().item() @@ -1116,6 +1119,7 @@ class SiglipFlashAttention2(SiglipAttention): past_key_value: tuple[torch.Tensor] | None = None, output_attentions: bool = False, use_cache: bool = False, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: output_attentions = False @@ -1178,7 +1182,13 @@ class SiglipFlashAttention2(SiglipAttention): value_states = value_states.to(target_dtype) attn_output = self._flash_attention_forward( - query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + unpadding_metadata=unpadding_metadata, ) attn_output = attn_output.reshape(bsz, q_len, self.embed_dim).contiguous() @@ -1190,7 +1200,15 @@ class SiglipFlashAttention2(SiglipAttention): return attn_output, attn_weights def _flash_attention_forward( - self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + dropout=0.0, + softmax_scale=None, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, ): """ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token @@ -1218,7 +1236,12 @@ class SiglipFlashAttention2(SiglipAttention): if attention_mask is not None: batch_size = query_states.shape[0] query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( - query_states, key_states, value_states, attention_mask, query_length + query_states, + key_states, + value_states, + attention_mask, + query_length, + unpadding_metadata=unpadding_metadata, ) cu_seqlens_q, cu_seqlens_k = cu_seq_lens @@ -1245,8 +1268,18 @@ class SiglipFlashAttention2(SiglipAttention): return attn_output - def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): - indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + def _upad_input( + self, + query_layer, + key_layer, + value_layer, + attention_mask, + query_length, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, + ): + if unpadding_metadata is None: + unpadding_metadata = _get_unpad_data(attention_mask) + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = unpadding_metadata batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape key_layer = index_first_axis( @@ -1316,6 +1349,7 @@ class SiglipEncoderLayer(nn.Module): hidden_states: torch.Tensor, attention_mask: torch.Tensor, output_attentions: bool | None = False, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, ) -> tuple[torch.FloatTensor]: """ Args: @@ -1330,11 +1364,19 @@ class SiglipEncoderLayer(nn.Module): residual = hidden_states hidden_states = self.layer_norm1(hidden_states) - hidden_states, attn_weights = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - output_attentions=output_attentions, - ) + if self._use_flash_attention_2: + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + unpadding_metadata=unpadding_metadata, + ) + else: + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) hidden_states = residual + hidden_states residual = hidden_states @@ -1435,6 +1477,7 @@ class SiglipEncoder(nn.Module): self.config = config self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]) self.gradient_checkpointing = False + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" # Ignore copy def forward( @@ -1475,6 +1518,9 @@ class SiglipEncoder(nn.Module): all_attentions = () if output_attentions else None hidden_states = inputs_embeds + unpadding_metadata = ( + _get_unpad_data(attention_mask) if self._use_flash_attention_2 and attention_mask is not None else None + ) for encoder_layer in self.layers: if output_hidden_states: encoder_states = encoder_states + (hidden_states,) @@ -1484,12 +1530,14 @@ class SiglipEncoder(nn.Module): hidden_states, attention_mask, output_attentions, + unpadding_metadata, ) else: layer_outputs = encoder_layer( hidden_states, attention_mask, output_attentions=output_attentions, + unpadding_metadata=unpadding_metadata, ) hidden_states = layer_outputs[0] @@ -4255,15 +4303,21 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S torch.Tensor: mask """ - ret = torch.zeros(size, size, device=device, dtype=torch.bool) - for i in range(size): - if num_left_chunks < 0: - start = 0 - else: - start = max((i // chunk_size - num_left_chunks) * chunk_size, 0) - ending = min((i // chunk_size + 1) * chunk_size + num_lookhead, size) - ret[i, start:ending] = True - return ret + indices = torch.arange(size, device=device) + chunk_index = indices // chunk_size + chunk_end = torch.clamp( + (chunk_index + 1) * chunk_size + num_lookhead, + max=size, + ) + + mask = indices.unsqueeze(0) < chunk_end.unsqueeze(1) + if num_left_chunks >= 0: + chunk_start = torch.clamp( + (chunk_index - num_left_chunks) * chunk_size, + min=0, + ) + mask.logical_and_(indices.unsqueeze(0) >= chunk_start.unsqueeze(1)) + return mask def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): input_lengths_after_cnn = (input_lengths - 1) // 2 + 1 diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 642d8ec02..4551126a7 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -134,6 +134,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): try: importlib.import_module("kernels") diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 99dafa2f2..fdd31ad21 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -127,6 +127,29 @@ class OmniPlatform(Platform): """ raise NotImplementedError + @classmethod + def validate_diffusion_attn_backend(cls, selected_backend: str) -> None: + """Reject an explicitly selected backend this platform cannot run. + + Platforms call this from ``get_diffusion_attn_backend_cls`` so that a + backend restricted to other hardware, or one whose kernel package is + missing, fails during resolution rather than at the first forward. + """ + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + backend_upper = selected_backend.upper() + backend_cls = DiffusionAttentionBackendEnum[backend_upper].get_class() + + supported = backend_cls.supported_platforms + platform = cls._omni_enum.value + if supported is not None and platform not in supported: + raise ValueError( + f"The {backend_upper} diffusion attention backend runs on {', '.join(supported)} only, " + f"but the current platform is {platform}. Select a backend supported here, " + "such as FLASH_ATTN or TORCH_SDPA." + ) + backend_cls.validate_available() + @classmethod def supports_torch_inductor(cls) -> bool: """Check if the platform supports torch.compile with inductor backend.""" diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 0774bca3a..02ce155ee 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -80,6 +80,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index bb0e1b436..0421814a0 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -130,6 +130,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 6ae65ed2b..630777058 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -102,6 +102,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index aa26cd82c..22cebeead 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -53,6 +53,7 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/quantization/int8_config.py b/vllm_omni/quantization/int8_config.py index 37d430047..ea096f9d5 100644 --- a/vllm_omni/quantization/int8_config.py +++ b/vllm_omni/quantization/int8_config.py @@ -50,9 +50,55 @@ if TYPE_CHECKING: # Dynamic quantization is supported first. ACTIVATION_SCHEMES = ["dynamic"] +# Ascend's npu_quant_matmul (QuantBatchMatmulV3) refuses a weight whose last +# dimension exceeds this, and it fails at the first forward rather than at load +# time. Layers still wider than the limit after TP sharding stay unquantized. +NPU_QUANT_MATMUL_MAX_OUT_FEATURES = 65535 + logger = init_logger(__name__) +def _fell_back_to_unquantized_npu( + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, +) -> bool: + """Swap a layer to unquantized weights when npu_quant_matmul cannot run its shape. + + The check uses the per-partition output size because that is what the kernel + actually sees; a layer over the limit on one rank can be within it at a + higher TP degree. Returns True when the layer was swapped and its weights + created, in which case the caller must not create its own. + """ + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition <= NPU_QUANT_MATMUL_MAX_OUT_FEATURES: + return False + + logger.warning_once( + "Keeping a %d-wide linear unquantized: npu_quant_matmul rejects an output dimension past " + "%d. Tensor parallelism shrinks this per-rank dimension, so a higher TP degree brings such " + "layers back into range.", + output_size_per_partition, + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + ) + fallback = UnquantizedLinearMethod() + layer.quant_method = fallback + fallback.create_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ) + return True + + def create_weight_parameter( output_size_per_partition: int, input_size_per_partition: int, @@ -231,6 +277,21 @@ class LazyWeightMixin: uses_meta_device: bool = True + # This mixin knows when a layer's weight is final, so it can hand the layer + # back to the host right there. Loaders that intend to offload the whole + # model after loading opt in per layer via ``enable_offload_after_quant``. + supports_offload_after_quant: bool = True + _offload_after_quant: bool = False + + def enable_offload_after_quant(self) -> None: + """Return each layer to host memory as soon as it has been quantized. + + Caps the load-time device footprint at one layer instead of the whole + model. A quant method instance belongs to a single layer, so this is not + a global switch. + """ + self._offload_after_quant = True + def create_weights( self, layer: torch.nn.Module, @@ -287,6 +348,11 @@ class LazyWeightMixin: # anything layer._already_called_process_weights_after_loading = True + # This layer's weight is final, so nothing needs it on the + # accelerator until inference. + if self._offload_after_quant: + layer.to("cpu") + # Note that we keep `layer._loaded_numel` around just in case # there is logic added to vllm in the future which calls a # weight loader twice - we do not want to re-initialize in @@ -354,6 +420,11 @@ class NPUInt8LinearMethod(BaseInt8LinearMethod): def __init__(self, quant_config: DiffusionInt8Config): super().__init__(quant_config) + def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None: + if _fell_back_to_unquantized_npu(layer, *args, **kwargs): + return + super().create_weights(layer, *args, **kwargs) + def process_weights_after_loading(self, layer: Module) -> None: layer.weight.data = layer.weight.data.t().contiguous() layer.weight_scale.data = layer.weight_scale.data.squeeze() @@ -424,6 +495,13 @@ class NPUInt8OnlineLinearMethod(LazyWeightMixin, NPUInt8LinearMethod): and quantized the weights during loading. """ + def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None: + # NPUInt8LinearMethod's override is unreachable from here: LazyWeightMixin + # comes first in the MRO and does not call super(). + if _fell_back_to_unquantized_npu(layer, *args, **kwargs): + return + super().create_weights(layer, *args, **kwargs) + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return