diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index c05c52732..c435f311e 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -37,24 +37,7 @@ steps: timeout_in_minutes: 20 depends_on: image-build commands: - - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py - agents: - queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - environment: - - "HF_HOME=/fsx/hf_cache" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - - label: "Diffusion Images API LoRA E2E" - timeout_in_minutes: 20 - depends_on: image-build - commands: - - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py + - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model" agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: @@ -137,24 +120,6 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Diffusion Tensor Parallelism Test" - timeout_in_minutes: 20 - depends_on: image-build - commands: - - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py - agents: - queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Diffusion GPU Worker Test" timeout_in_minutes: 20 depends_on: image-build diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 8e380ad6d..fc9fea25c 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -29,7 +29,7 @@ steps: timeout_in_minutes: 20 depends_on: image-build commands: - - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py + - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model" agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 2466bc71e..c67a9c7f4 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -12,8 +12,8 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT Essential Elements of an Effective PR Description Checklist - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". -- [ ] The test plan, such as providing test command. -- [ ] The test results, such as pasting the results comparison before and after, or e2e results +- [ ] The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/) +- [ ] The test results. Please pasting the results comparison before and after, or e2e results. - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.** - [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft. diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 05b37a661..d2982fedb 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,5 +1,5 @@ ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.15.0 +ARG VLLM_BASE_TAG=v0.16.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} @@ -14,7 +14,7 @@ RUN apt-get update && \ # Install vLLM wheel for the specific commit # URL format: https://wheels.vllm.ai// -# Current vLLM commit: f97ca671766c5201404e9fc812e35bf2c4e95a01 +# Current vLLM wheel commit: ec12d39d44739bee408ec1473acc09e75daf1a5d # Note: --prerelease=allow is required because nightly/commit wheels use # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default, # causing it to fall back to the older stable release from PyPI. diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index ab4cdd5c1..577ca90de 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md new file mode 100644 index 000000000..dbf87d352 --- /dev/null +++ b/docs/contributing/ci/CI_5levels.md @@ -0,0 +1,678 @@ +# Multi-Level Automated Testing System Documentation + +## Document Overview + +This testing system aims to build a complete, efficient, and well-structured quality assurance framework for the development, integration, and release of model services. It draws on the concept of the test pyramid from modern software engineering, progressively expanding testing activities from basic code logic verification to complex end-to-end (E2E) functionality, performance, accuracy, and even long-term stability validation. + +Through five levels (L1-L5) and common (Common) specifications, the system clarifies the testing objectives, scope, execution frequency, and required resources for different development stages (e.g., each commit, PR merge, daily build, pre-release). This ensures that models meet high standards for functionality, performance, and reliability across various deployment scenarios (online serving and offline inference). + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
LevelScope & FocusTime CostTest DirDocFrequencyHardware
CommonContribution Guideline & PR checklist//.github/PULL_REQUEST_TEMPLATE.md Test Style (PR Checklist)//
CI Failure Description// CI Failures//
L1
(Unit & Logic)		Unit tests for components like entrypoints, models<15min/tests/{component_name}/test_xxx + Chapter 1
+ Section 1 L1&L2: Purpose, Test Content, Directory Location, Example + 		PR with ready label (also can run locally)CPU
L2
(E2E across models & GPU-required UT)		Online & Offline (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT		 + /tests/e2e/online_serving/test_{model_name}.py
+ /tests/e2e/offline_inference/test_{model_name}.py + 		+ Chapter 1
+ L1&L2: Purpose, Test Content, Directory Location, Example + 		PR with ready labelGPU
L3
(Important Perf & Integration & Accuracy)		Online & Offline (multiple deployment scenarios):
real model, normal inference function, normal accuracy		<30min + /tests/e2e/online_serving/test_{model_name}_expansion.py
+ /tests/e2e/offline_inference/test_{model_name}_expansion.py + 		+ Chapter 2
+ L3: Purpose, Test Content, Directory Location, Example + 		PR Merged (Also run L1&L2 Tests)GPU
L4
(Perf & Integration & Accuracy)		Online & Offline: full functional scenarios + performance test + doc test<3 hour + Full Function:
+ /tests/e2e/online_serving/test_{model_name}_expansion.py
+ /tests/e2e/offline_inference/test_{model_name}_expansion.py
+ Performance:
+ /tests/e2e/perf/nightly.json
+ Doc Test:
+ tests/example/online_serving/test_{model_name}.py
+ tests/example/offline_inference/test_{model_name}.py + 		+ Chapter 3
+ L4: Purpose, Test Content, Directory Location, Example + 		NightlyGPU
L5
(Stability & Reliability)		Online & Offline: long-term stability test + reliability test Depends on reality + Stability:
+ tests/e2e/stability/weekly.json
+ Reliability:
+ tests/e2e/reliability/test_{model_name}.py + 		+ Chapter 4
+ L5: Purpose, Test Content, Directory Location, Example + 		Weekly / Days before ReleaseGPU
+ + +--- +
+ The folder structure for tests file based on the 5 levels design +Legend: `✅` = test exists, `⬜` = suggested to add. +``` +vllm_omni/ tests/ +├── config/ → ├── config/ +│ ├── model.py │ └── test_model.py ⬜ +│ └── lora.py │ └── test_lora.py ⬜ +│ +├── core/ → ├── core/ +│ └── sched/ │ └── sched/ +│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py ⬜ +│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py ⬜ +│ └── output.py │ └── test_output.py ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData) +│ +├── diffusion/ → ├── diffusion/ +│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py ⬜ +│ ├── attention/ │ ├── attention/ +│ │ ├── layer.py │ │ ├── test_attention_sp.py ✅ +│ │ └── backends/ │ │ └── test_flash_attn.py ✅ +│ ├── distributed/ │ ├── distributed/ +│ │ └── ... │ │ ├── test_comm.py ✅ +│ │ │ │ ├── test_cfg_parallel.py ✅ +│ │ │ │ └── test_sp_plan_hooks.py ✅ +│ ├── lora/ │ ├── lora/ +│ │ └── ... │ │ ├── test_base_linear.py ✅ +│ │ │ │ └── test_lora_manager.py ✅ +│ ├── models/ │ ├── models/ +│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage) +│ │ ├── z_image/ │ │ └── z_image/ +│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅ +│ └── worker/ │ └── worker/ +│ ├── diffusion_worker.py │ └── test_diffusion_worker.py ✅ file at diffusion/test_diffusion_worker.py +│ └── diffusion_model_runner.py │ +│ +├── distributed/ → ├── distributed/ +│ └── omni_connectors/ │ └── omni_connectors/ +│ ├── adapter.py │ ├── test_adapter_and_flow.py ✅ +│ ├── kv_transfer_manager.py │ ├── test_basic_connectors.py ✅ +│ ├── connectors/ │ ├── test_kv_flow.py ✅ +│ └── utils/ │ └── test_omni_connector_configs.py ✅ +│ +├── engine/ → ├── engine/ +│ ├── input_processor.py │ ├── test_input_processor.py ⬜ (no processor.py in source) +│ ├── output_processor.py │ └── test_output_processor.py ⬜ +│ └── arg_utils.py │ └── test_arg_utils.py ⬜ +│ +├── entrypoints/ → ├── entrypoints/ +│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online) +│ ├── omni_llm.py │ ├── test_omni_llm.py ✅ +│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py) +│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅ +│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py +│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅ +│ ├── stage_utils.py │ ├── test_stage_utils.py ✅ +│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve) +│ │ └── ... │ │ └── test_*.py ⬜ +│ └── openai/ │ └── openai_api/ # maps to entrypoints/openai/ +│ ├── api_server.py │ ├── test_api_server.py ⬜ (e2e indirect coverage) +│ ├── serving_chat.py │ ├── test_serving_chat_sampling_params.py ✅ +│ ├── serving_speech.py │ ├── test_serving_speech.py ✅ +│ └── image_api_utils.py │ └── test_image_server.py ✅ +│ +├── inputs/ → ├── inputs/ +│ ├── data.py │ ├── test_data.py ⬜ +│ ├── parse.py │ ├── test_parse.py ⬜ +│ └── preprocess.py │ └── test_preprocess.py ✅ currently in entrypoints/test_omni_input_preprocessor.py +│ +├── model_executor/ → ├── model_executor/ +│ ├── layers/ │ ├── layers/ +│ │ └── mrope.py │ │ └── test_mrope.py ⬜ +│ ├── model_loader/ │ ├── model_loader/ +│ │ └── weight_utils.py │ │ └── test_weight_utils.py ⬜ +│ ├── models/ │ ├── models/ +│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/ +│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_audio_length.py ✅ +│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_thinker.py ⬜ +│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ ├── test_qwen2_5_omni_talker.py ⬜ +│ │ └── qwen3_omni/ │ │ │ └── test_qwen2_5_omni_token2wav.py ⬜ +│ │ └── ... │ │ └── qwen3_omni/ +│ ├── stage_configs/ │ │ └── test_*.py ⬜ +│ │ └── *.yaml │ └── stage_configs/ (used by e2e, test_*.py can be added) ⬜ +│ └── stage_input_processors/ │ └── stage_input_processors/ +│ └── ... │ └── test_*.py ⬜ +│ +├── sample/ → ├── sample/ +│ └── __init__.py │ └── test_*.py ⬜ +│ +├── utils/ → ├── utils/ +│ └── __init__.py │ └── test_*.py ⬜ (no platform_utils.py currently) +│ +├── worker/ → ├── worker/ +│ ├── gpu_ar_model_runner.py │ ├── test_gpu_ar_model_runner.py ⬜ +│ ├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py ⬜ +│ ├── gpu_generation_model_runner.py │ ├── test_gpu_generation_model_runner.py ✅ +│ ├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py ⬜ +│ ├── gpu_model_runner.py │ ├── test_omni_gpu_model_runner.py ✅ +│ └── mixins.py │ └── (npu under platforms/npu/worker/) # not worker/npu/ +│ +├── platforms/ → (no tests/platforms/, e2e and stage_configs provide indirect coverage) +│ ├── cuda/ +│ ├── npu/worker/ # NPU worker here, not vllm_omni/worker/npu/ +│ ├── rocm/ +│ └── xpu/worker/ +│ +├── outputs.py → test_outputs.py ✅ (at tests root) +├── (logger, patch, request, version) → (no corresponding unit test) +│ +└── e2e (tests side only) → ├── e2e/ + ├── online_serving/ ✅ non-empty + │ ├── test_async_omni.py + │ ├── test_qwen3_omni.py + │ ├── test_qwen3_omni_expansion.py + │ ├── test_image_gen_edit.py + │ ├── test_images_generations_lora.py + │ └── stage_configs/ + └── offline_inference/ ✅ + ├── test_qwen2_5_omni.py + ├── test_qwen3_omni.py + ├── test_bagel_text2img.py + ├── test_t2i_model.py + ├── test_t2v_model.py + ├── test_ovis_image.py + ├── test_zimage_tensor_parallel.py + ├── test_cache_dit.py + ├── test_teacache.py + ├── test_stable_audio_model.py + ├── test_diffusion_cpu_offload.py + ├── test_diffusion_layerwise_offload.py + ├── test_diffusion_lora.py + ├── test_sequence_parallel.py + └── stage_configs/ + ├── qwen2_5_omni_ci.yaml + ├── qwen3_omni_ci.yaml + ├── bagel_*.yaml + └── npu/, rocm/, etc. +``` + + +
+ + +## Common Specifications + +Before entering specific testing levels, the project establishes two common specifications aimed at standardizing the development process and quickly locating issues. + +1. ***PR Checklist ([Tests Style](../ci/tests_style.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance. +2. ***CI Failure Explanation ([CI Failures](../ci/failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency. + +## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification + +### 1.1 Testing Purpose + +L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed. + +L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle . + + + +### 1.2 Testing Content and Scope + +- ***L1 (Unit & Logic Testing)***: +- - ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`. + - ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack. + - ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback. +- ***L2 (Basic End-to-End Testing)***: +- - ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference). + - ***Focus***: Uses `dummy` models or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, etc. Also includes some unit tests that require launching independent service instances. + - ***Characteristic***: Requires ***GPU*** resources to perform model computations. + +### 1.3 Test Directory and Execution Files + +A clear directory structure is key to managing test cases efficiently. + +- ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py` +- - Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file. +- ***L2 Test Directory***: +- - Online Serving: `/tests/e2e/online_serving/test_{model_name}.py` + - Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py` + +### 1.4 Execution Method and Example + +- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered. +- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment. +- ***Script Example***: + +
+ L1 Test Examples + +Examples from `tests/model_executor/models/qwen2_5_omni/test_audio_length.py` +```python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +def test_resolve_max_mel_frames_default(): + from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames + + assert resolve_max_mel_frames(None, default=30000) == 30000 + assert resolve_max_mel_frames(None, default=6000) == 6000 + + +def test_resolve_max_mel_frames_explicit(): + from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames + + # Explicit argument always wins over default + assert resolve_max_mel_frames(123, default=30000) == 123 + assert resolve_max_mel_frames(6000, default=30000) == 6000 + assert resolve_max_mel_frames(0, default=30000) == 0 + + +@pytest.mark.parametrize("repeats", [2, 4]) +@pytest.mark.parametrize("code_len", [0, 1, 32768]) +@pytest.mark.parametrize("max_mel_frames", [None, -1, 0, 1, 6000, 30000]) +def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames): + """Guard that any max_mel_frames yields a mel length aligned to repeats, and + consistent with the truncated code length (prevents concat mismatch). + """ + from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length + + target_code_len, target_mel_len = cap_and_align_mel_length( + code_len=code_len, + repeats=repeats, + max_mel_frames=max_mel_frames, + ) + + assert isinstance(target_code_len, int) + assert isinstance(target_mel_len, int) + + if code_len == 0: + assert target_code_len == 0 + assert target_mel_len == 0 + return + + assert target_code_len >= 1 + assert target_mel_len >= repeats + assert target_mel_len % repeats == 0 + assert target_mel_len == target_code_len * repeats + assert target_code_len <= code_len + + if max_mel_frames is not None and int(max_mel_frames) > 0 and int(max_mel_frames) >= repeats: + assert target_mel_len <= int(max_mel_frames) +``` +
+ +
+ L2 Test Examples +You can refer to Test Examples in Chapter 2 to see example test cases that incorporate both L2 and L3 testing logic. +
+ +- - ***Run Command***: + + `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py` + `pytest -s -v -m 'core_model and cpu' --run-level=core_model` + +## Chapter 2: L3 Level Testing - Core Integration, Performance, and Accuracy Verification + +### 2.1 Testing Purpose + +L3 level testing executes after code is merged into the main branch. Its core purpose is to verify the integration effect, key performance indicators, and output accuracy of ***real models*** in ***multiple deployment scenarios*** + +. It acts as the "quality gatekeeper" for the main branch, ensuring that no merge breaks the core capabilities of the model service. Testing needs to provide clear conclusions within a relatively short time (<30min), balancing test depth with feedback speed. + + + +### 2.2 Testing Content and Scope + +- ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc. +- ***Core Verification***: +- 1. ***Inference Functionality***: Ensures real models can perform forward computation normally and return results. + 2. ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues. + 3. ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds. + +### 2.3 Test Directory and Execution Files + +- ***Functional Testing***: +- - Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py` + - Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py` + - (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests). + +### 2.4 Execution Method and Example + +- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch. +- ***Execution Environment***: ***GPU*** servers. +- ***Script Example***: + +???+ example "Test Examples" + + **2.4.1 Mark Declaration Section** + + ```python + @pytest.mark.advanced_model + @pytest.mark.core_model + @pytest.mark.parametrize("omni_server", test_params, indirect=True) + ``` + + **Explanation**: + + @pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level. + + @pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly. + + @pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing. + + **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model. + + **2.4.2 Test Function Definition and Documentation** + + ```python + def test_mix_to_text_audio_001(omni_server, openai_client) -> None: + """ + Test multi-modal input processing and text/audio output generation via OpenAI API. + Deploy Setting: default yaml + Input Modal: text + audio + video + image + Output Modal: text + audio + Input Setting: stream=True + Datasets: single request + """ + ``` + + **Explanation**: + + **Function Naming Convention**: Uses the test_ prefix, describes the test scenario mix_to_text_audio, and the number 001 indicates the first test case for this scenario. + + **Parameter Explanation**: + + omni_server: Omni server instance obtained via fixture, containing model information and configuration. + + openai_client: Unified OpenAI client processing instance, encapsulating request sending and response validation logic. + + Docstring: Describes the test purpose, deployment settings, input/output modalities, streaming settings, and dataset type in detail, providing clear context for test maintenance. + + **2.4.3 Multimodal Data Preparation** + + ```python + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" + ``` + + **Explanation**: + + **Data Generation Functions**: Use the generate_synthetic_* series of functions to generate synthetic test data, avoiding reliance on external resources and ensuring test reproducibility and stability. + + **Parameter Explanation**: + + Video: width, height, duration_frames + + Image: width, height + + Audio: duration_seconds, channels + + **2.4.4 Request Configuration and Keyword Validation** + + ```python + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "key_words": { + "audio": ["water", "cricket"], + "video": ["sphere", "globe", "circle", "round"], + "image": ["square", "quadrate"], + "text": ["beijing"] + }, + } + ``` + + **Explanation**: + + **Model Specification**: Uses omni_server.model to ensure the test aligns with the model configured on the server. + + **Keyword Validation Mechanism**: This is an innovative design of the template to address the specific needs of multimodal testing: + + Audio Keywords: Validate whether the generated text's description of audio content contains expected elements (e.g., "water" for water sounds, "cricket" for cricket sounds). If you provide multiple keywords, the validation is considered successful if at least one keyword is present. + + **Video Keywords**: Validate whether the generated text's description of video content contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present. + + Image Keywords: Validate whether the generated text's description of image content contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present. + + Text Keywords: Validate whether the generated text contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present. + + **2.4.5 Request Execution** + + ```python + openai_client.send_request(request_config, request_num=1) + ``` + + **Explanation**: + + **Unified Client**: Uses the OpenAIClientHandler instance to send requests. This client encapsulates error handling, retry mechanisms, and response validation logic. + + **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n. + + **Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation. + +- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` + +## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing + +### 3.1 Testing Purpose + +L4 level testing is a comprehensive quality audit before a version release. It expands upon L3, executing ***full*** functional scenarios, conducting systematic ***performance stress tests***, and simultaneously verifying the correctness of accompanying ***example documentation***. Its purpose is to perform deep validation of the system during off-peak nighttime hours, providing quality trend reports for daytime development and data support for release decisions. + + + +### 3.2 Testing Content and Scope + +- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling. +- ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization. +- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description. + +### 3.3 Test Directory and Execution Files + +- ***Functional Testing***: Same directories as L3. +- ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json` +- ***Documentation Example Tests***: +- - `tests/example/online_serving/test_{model_name}.py` + - `tests/example/offline_inference/test_{model_name}.py` + +### 3.4 Execution Method and Example + +- ***Trigger Timing***: **`Nightly`**, automatically executed every night. +- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing. +- ***Script Example***: +
+ Test Examples +When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json: + +```JSON +{ + "test_name": "test_qwen3_omni", + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "stage_config_name": "qwen3_omni.yaml" + }, + "benchmark_params": [ + { + "dataset_name": "random", + "num_prompts": [10, 20], + "request_rate": [0.5, 1], + "random_input_len": 2500, + "random_output_len": 900, + "ignore_eos": true, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": 100000, + "mean_audio_ttfp_ms": 100000, + "mean_audio_rtf": 100000 + } + } + ] +} +``` + +#### Parameter Explanation + +***Overview*** + +| Field | Required | Description | +| ---------------- | -------- | --------------------------------------------------------------- | +| test_name | Yes | Unique identifier for the test case | +| server_params | Yes | Server-side configuration parameters | +| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) | + +#### server_params Configuration + +##### Basic Parameters + +| Parameter | Required | Example | Description | +| ----------------- | -------- | ---------------------------------- | ----------------------------- | +| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path | +| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name | + +##### Dynamic Configuration (update/delete) + +Supports incremental modifications based on the basic configuration: + +| Operation | Description | +| --------- | ------------------------------------ | +| update | Update or add configuration items | +| delete | Delete specified configuration items | + +***Example***: +``` +"update": { + "async_chunk": true, // Enable asynchronous chunk processing + "stage_args": { + "0": { + "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" + } + } +}, +"delete": { + "stage_args": { + "2": ["custom_process_input_func"] // Delete this configuration for stage 2 + } +} +``` +#### benchmark_params Configuration + +You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows: + +1. Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file. +2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file. +3. Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark. +4. The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below: + +| Parameter | Type | Required | Example/Values | Description | +| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| 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 | int / array | No | 1, [1, 2, 3] | 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] | 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. | +
+ +- - ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`). + +## Chapter 4: L5 Level Testing - Stability and Reliability Testing + +### 4.1 Testing Purpose + +L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness. + + + +### 4.2 Testing Content and Scope + +- ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable. +- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities. + +### 4.3 Test Directory and Execution Files + +- ***Stability Test Configuration***: `tests/e2e/stability/weekly.json` +- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py` + +### 4.4 Execution Method and Example + +- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower. +- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment. +- ***Script Example***: +
+ Test Examples +```python +# WIP +``` +
+ +- - ***Stability***: (Execution would be driven by the configuration in `weekly.json`). + - ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py` + +## Summary + +This multi-level testing system achieves continuous, progressive validation of model service quality by tightly integrating testing activities with the development workflow (commit, review, merge, release). From rapid unit testing to comprehensive end-to-end testing, and further to in-depth performance, stability, and reliability verification, each level has clear objectives, collectively building a robust quality protection net. By following this system, teams can deliver high-quality, highly reliable model services more efficiently. diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md index d64d98a60..6e3c10ac8 100644 --- a/docs/contributing/ci/failures.md +++ b/docs/contributing/ci/failures.md @@ -1,4 +1,10 @@ # CI Failures -What should I do when a CI job fails on my PR, but I don't think my PR caused -the failure? +What should I do when a CI job fails on my PR, but I don't think my PR caused the failure? + + +## Common Case 1: + +🚨 Error: The command was interrupted by a signal: signal: terminated + +Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass. diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index 2afff8460..c9e5d18e6 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -7,7 +7,8 @@ Defined in `pyproject.toml`: | Marker | Description | | ------------------ | --------------------------------------------------------- | -| `core_model` | Core model tests (run in each PR) | +| `core_model` | L1&L2 tests (run in each PR) | +| `advanced_model` | L3&L4 level tests (run in each merge or nightly) | | `diffusion` | Diffusion model tests | | `omni` | Omni model tests | | `cache` | Cache backend tests | diff --git a/docs/contributing/ci/tests_style.md b/docs/contributing/ci/tests_style.md index 65c2b0443..ab94f2f56 100644 --- a/docs/contributing/ci/tests_style.md +++ b/docs/contributing/ci/tests_style.md @@ -2,8 +2,18 @@ To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. This document outlines the guidelines for organizing and naming test files. +## Checklist before submitting your test files + +1. The file is saved in an appropriate place and the file name is clear. +2. The coding style follows the requirements outlined below. +3. All test functions have appropriate pytest markers. +4. For tests that need run in CI, please ensure it labeled as ``@pytest.mark.core_model` the test is configured under the `./buildkite/` folder. + + ## Test Types +For more details about our [Five Levels Tests design](../ci/CI_5levels.md). + ### Unit Tests and System Tests For unit tests and system tests, we strongly recommend placing test files in the same directory structure as the source code being tested, using the naming convention `test_*.py`. @@ -18,95 +28,141 @@ End-to-end tests verify the complete functionality of a system or component. For ## Test Directory Structure -The ideal directory structure mirrors the source code organization: +The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add. ``` -vllm_omni/ tests/ -├── config/ → ├── config/ -│ └── model.py │ └── test_model.py +vllm_omni/ tests/ +├── config/ → ├── config/ +│ ├── model.py │ └── test_model.py ⬜ +│ └── lora.py │ └── test_lora.py ⬜ │ -├── core/ → ├── core/ -│ └── sched/ │ └── sched/ # Maps to core/sched/ -│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py -│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py -│ └── output.py │ └── test_output.py +├── core/ → ├── core/ +│ └── sched/ │ └── sched/ +│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py ⬜ +│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py ⬜ +│ └── output.py │ └── test_output.py ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData) │ -├── diffusion/ → ├── diffusion/ -│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py -│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py -│ ├── attention/ │ ├── attention/ # Maps to diffusion/attention/ -│ │ └── backends/ │ │ └── test_*.py -│ ├── models/ │ ├── models/ # Maps to diffusion/models/ -│ │ ├── qwen_image/ │ │ ├── qwen_image/ -│ │ │ └── ... │ │ │ └── test_*.py -│ │ └── z_image/ │ │ └── z_image/ -│ │ └── ... │ │ └── test_*.py -│ └── worker/ │ └── worker/ # Maps to diffusion/worker/ -│ └── ... │ └── test_*.py +├── diffusion/ → ├── diffusion/ +│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py ⬜ +│ ├── attention/ │ ├── attention/ +│ │ ├── layer.py │ │ ├── test_attention_sp.py ✅ +│ │ └── backends/ │ │ └── test_flash_attn.py ✅ +│ ├── distributed/ │ ├── distributed/ +│ │ └── ... │ │ ├── test_comm.py ✅ +│ │ │ │ ├── test_cfg_parallel.py ✅ +│ │ │ │ └── test_sp_plan_hooks.py ✅ +│ ├── lora/ │ ├── lora/ +│ │ └── ... │ │ ├── test_base_linear.py ✅ +│ │ │ │ └── test_lora_manager.py ✅ +│ ├── models/ │ ├── models/ +│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage) +│ │ ├── z_image/ │ │ └── z_image/ +│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅ +│ └── worker/ │ └── worker/ +│ ├── diffusion_worker.py │ └── test_diffusion_worker.py ✅ file at diffusion/test_diffusion_worker.py +│ └── diffusion_model_runner.py │ │ -├── distributed/ → ├── distributed/ -│ └── ... │ └── test_*.py +├── distributed/ → ├── distributed/ +│ └── omni_connectors/ │ └── omni_connectors/ +│ ├── adapter.py │ ├── test_adapter_and_flow.py ✅ +│ ├── kv_transfer_manager.py │ ├── test_basic_connectors.py ✅ +│ ├── connectors/ │ ├── test_kv_flow.py ✅ +│ └── utils/ │ └── test_omni_connector_configs.py ✅ │ -├── engine/ → ├── engine/ -│ ├── processor.py │ ├── test_processor.py -│ └── output_processor.py │ └── test_output_processor.py +├── engine/ → ├── engine/ +│ ├── input_processor.py │ ├── test_input_processor.py ⬜ (no processor.py in source) +│ ├── output_processor.py │ └── test_output_processor.py ⬜ +│ └── arg_utils.py │ └── test_arg_utils.py ⬜ │ -├── entrypoints/ → ├── entrypoints/ -│ ├── omni_llm.py │ ├── test_omni_llm.py # UT: OmniLLM core logic (mocked) -│ ├── omni_stage.py │ ├── test_omni_stage.py # UT: OmniStage logic -│ ├── omni.py │ ├── test_omni.py # E2E: Omni class (offline inference) -│ ├── async_omni.py │ ├── test_async_omni.py # E2E: AsyncOmni class -│ ├── cli/ │ ├── cli/ # Maps to entrypoints/cli/ -│ │ └── ... │ │ └── test_*.py -│ └── openai/ │ └── openai/ # Maps to entrypoints/openai/ -│ ├── api_server.py │ ├── test_api_server.py # E2E: API server (online serving) -│ └── serving_chat.py │ └── test_serving_chat.py +├── entrypoints/ → ├── entrypoints/ +│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online) +│ ├── omni_llm.py │ ├── test_omni_llm.py ✅ +│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py) +│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅ +│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py +│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅ +│ ├── stage_utils.py │ ├── test_stage_utils.py ✅ +│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve) +│ │ └── ... │ │ └── test_*.py ⬜ +│ └── openai/ │ └── openai_api/ # maps to entrypoints/openai/ +│ ├── api_server.py │ ├── test_api_server.py ⬜ (e2e indirect coverage) +│ ├── serving_chat.py │ ├── test_serving_chat_sampling_params.py ✅ +│ ├── serving_speech.py │ ├── test_serving_speech.py ✅ +│ └── image_api_utils.py │ └── test_image_server.py ✅ │ -├── inputs/ → ├── inputs/ -│ ├── data.py │ ├── test_data.py -│ ├── parse.py │ ├── test_parse.py -│ └── preprocess.py │ └── test_preprocess.py +├── inputs/ → ├── inputs/ +│ ├── data.py │ ├── test_data.py ⬜ +│ ├── parse.py │ ├── test_parse.py ⬜ +│ └── preprocess.py │ └── test_preprocess.py ✅ currently in entrypoints/test_omni_input_preprocessor.py │ -├── model_executor/ → ├── model_executor/ -│ ├── layers/ │ ├── layers/ -│ │ └── mrope.py │ │ └── test_mrope.py -│ ├── model_loader/ │ ├── model_loader/ -│ │ └── weight_utils.py │ │ └── test_weight_utils.py -│ ├── models/ │ ├── models/ -│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/ -│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_qwen2_5_omni_thinker.py # UT -│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_talker.py # UT -│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ └── test_qwen2_5_omni_token2wav.py # UT -│ │ └── qwen3_omni/ │ │ └── qwen3_omni/ -│ │ └── ... │ │ └── test_*.py -│ ├── stage_configs/ │ └── stage_configs/ # Configuration tests (if needed) -│ │ └── ... │ └── test_*.py -│ └── stage_input_processors/ │ └── stage_input_processors/ -│ └── ... │ └── test_*.py +├── model_executor/ → ├── model_executor/ +│ ├── layers/ │ ├── layers/ +│ │ └── mrope.py │ │ └── test_mrope.py ⬜ +│ ├── model_loader/ │ ├── model_loader/ +│ │ └── weight_utils.py │ │ └── test_weight_utils.py ⬜ +│ ├── models/ │ ├── models/ +│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/ +│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_audio_length.py ✅ +│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_thinker.py ⬜ +│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ ├── test_qwen2_5_omni_talker.py ⬜ +│ │ └── qwen3_omni/ │ │ │ └── test_qwen2_5_omni_token2wav.py ⬜ +│ │ └── ... │ │ └── qwen3_omni/ +│ ├── stage_configs/ │ │ └── test_*.py ⬜ +│ │ └── *.yaml │ └── stage_configs/ (used by e2e, test_*.py can be added) ⬜ +│ └── stage_input_processors/ │ └── stage_input_processors/ +│ └── ... │ └── test_*.py ⬜ │ -├── sample/ → ├── sample/ -│ └── ... │ └── test_*.py +├── sample/ → ├── sample/ +│ └── __init__.py │ └── test_*.py ⬜ │ -├── utils/ → ├── utils/ -│ └── platform_utils.py │ └── test_platform_utils.py +├── utils/ → ├── utils/ +│ └── __init__.py │ └── test_*.py ⬜ (no platform_utils.py currently) │ -├── worker/ → ├── worker/ - ├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py - ├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py - ├── gpu_model_runner.py │ ├── test_gpu_model_runner.py - └── npu/ │ └── npu/ # Maps to worker/npu/ - └── ... │ └── test_*.py +├── worker/ → ├── worker/ +│ ├── gpu_ar_model_runner.py │ ├── test_gpu_ar_model_runner.py ⬜ +│ ├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py ⬜ +│ ├── gpu_generation_model_runner.py │ ├── test_gpu_generation_model_runner.py ✅ +│ ├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py ⬜ +│ ├── gpu_model_runner.py │ ├── test_omni_gpu_model_runner.py ✅ +│ └── mixins.py │ └── (npu under platforms/npu/worker/) # not worker/npu/ │ -└── e2e/ → ├── e2e/ # End-to-end scenarios (no 1:1 source mirror) - ├── online_serving/ # Full-stack online serving flows - │ └── (empty for now) - └── offline_inference/ # Full offline inference flows - ├── test_qwen2_5_omni.py # Moved from multi_stages/ - ├── test_qwen3_omni.py # Moved from multi_stages_h100/ - ├── test_t2i_model.py # Moved from single_stage/ - └── stage_configs/ # Shared stage configs - ├── qwen2_5_omni_ci.yaml - └── qwen3_omni_ci.yaml +├── platforms/ → (no tests/platforms/, e2e and stage_configs provide indirect coverage) +│ ├── cuda/ +│ ├── npu/worker/ # NPU worker here, not vllm_omni/worker/npu/ +│ ├── rocm/ +│ └── xpu/worker/ +│ +├── outputs.py → test_outputs.py ✅ (at tests root) +├── (logger, patch, request, version) → (no corresponding unit test) +│ +└── e2e (tests side only) → ├── e2e/ + ├── online_serving/ ✅ non-empty + │ ├── test_async_omni.py + │ ├── test_qwen3_omni.py + │ ├── test_qwen3_omni_expansion.py + │ ├── test_image_gen_edit.py + │ ├── test_images_generations_lora.py + │ └── stage_configs/ + └── offline_inference/ ✅ + ├── test_qwen2_5_omni.py + ├── test_qwen3_omni.py + ├── test_bagel_text2img.py + ├── test_t2i_model.py + ├── test_t2v_model.py + ├── test_ovis_image.py + ├── test_zimage_tensor_parallel.py + ├── test_cache_dit.py + ├── test_teacache.py + ├── test_stable_audio_model.py + ├── test_diffusion_cpu_offload.py + ├── test_diffusion_layerwise_offload.py + ├── test_diffusion_lora.py + ├── test_sequence_parallel.py + └── stage_configs/ + ├── qwen2_5_omni_ci.yaml + ├── qwen3_omni_ci.yaml + ├── bagel_*.yaml + └── npu/, rocm/, etc. ``` @@ -117,6 +173,8 @@ vllm_omni/ tests/ - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py` +- **Expansion Tests** + ### Best Practices 1. **Mirror Source Structure**: Test directories should mirror the source code structure @@ -144,81 +202,217 @@ vllm_omni/ tests/ ### Template #### E2E - Online serving +E2E Online tests for Qwen3-Omni model with mix input and audio+text output. Based on `tests/e2e/online_serving/test_qwen3_omni.py`. + ```python -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Online E2E smoke test for an omni model (video,text,audio → audio). +E2E Online tests for Qwen3-Omni model with mix input and audio+text output. """ + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +import threading from pathlib import Path -import pytest import openai +import pytest -from tests.utils import hardware_test +from tests.conftest import ( + OmniServer, + convert_audio_to_text, + cosine_similarity_text, + dummy_messages_from_mix_data, + generate_synthetic_video, + merge_base64_and_convert_to_text, +) +from vllm_omni.platforms import current_omni_platform -# Optional: set process start method for workers -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +# Edit: model name and stage config path +models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] + +#If you use the default configuration file, you can directly use the following address. +def get_default_config(): + return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") + +#If you need to modify the configuration file, you can use modify_stage_config. +def get_chunk_config(): + path = modify_stage_config( + get_default_config(), + updates={ + "async_chunk": True, + "stage_args": { + 0: { + "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" + }, + 1: { + "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk" + }, + }, + }, + deletes={"stage_args": {2: ["custom_process_input_func"]}}, + ) + return path + +stage_configs = [get_default_config(), CHUNK_CONFIG_PATH] -models = ["{your model name}"] #Edit here to load your model -stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml test_params = [(model, stage_config) for model in models for stage_config in stage_configs] -#OmniServer，Used to start the vllm-omni server -class OmniServer: - xxx + +#Please use this method to launch the online instance. +_omni_server_lock = threading.Lock() + +@pytest.fixture(scope="module") +def omni_server(request): + """Start vLLM-Omni server as a subprocess. Use module scope for multi-stage init (10-20+ min).""" + with _omni_server_lock: + model, stage_config_path = request.param + with OmniServer( + model, + ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], + ) as server: + yield server @pytest.fixture -def omni_server(request): - model, stage_config_path = request.param - with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server: - yield server - - -#handle request message -@pytest.fixture(scope="session") -def base64_encoded_video() -> str: - xxx - -@pytest.fixture(scope="session") -def dummy_messages_from_video_data(video_data_url: str, content_text: str) -> str: - xxx - -@pytest.mark.core_model -@pytest.mark.omni -@hardware_test( - res={"cuda": "L4", "rocm": "MI325", "npu": "A2"}, - num_cards={"cuda": 2, "rocm": 2, "npu": 4}, -) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_video_to_audio( - client: openai.OpenAI, - omni_server, - base64_encoded_video: str, -) -> None: - #set message - video_data_url = f"data:video/mp4;base64, {base64_encoded_video}" - messages = dummy_messages_from_video_data(video_data_url) - - #send request - chat_completion = client.chat.completions.create( - model=omni_server.model, - messages=messages, +def client(omni_server): + """OpenAI client for the running vLLM-Omni server.""" + return openai.OpenAI( + base_url=f"http://{omni_server.host}:{omni_server.port}/v1", + api_key="EMPTY", ) - #verify text output - text_choice = chat_completion.choices[0] - assert text_choice.finish_reason == "length" +#Please use function definitions above the test function to define the prompts and other parameters you need. +def get_system_prompt(): + return { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual inputs, " + "as well as generating text and speech." + ), + } + ], + } - #verify audio output - audio_choice = chat_completion.choices[1] - audio_message = audio_choice.message - if hasattr(audio_message, "audio") and audio_message.audio: - assert audio_message.audio.data is not None - assert len(audio_message.audio.data) > 0 +... + +#Please define test case tags according to the instructions in the marker documentation. +@pytest.mark.core_model +@pytest.mark.omni +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None: + # PLEASE FOLLOW THESE TEMPLATE INSTRUCTIONS: + # ============================================================================ + # TEMPLATE USAGE GUIDE: + # 1. Copy this entire function as a starting point for multi-modal tests + # 2. Update the test name to reflect your specific test scenario + # 3. Modify input/output modalities as needed (see OPTIONS section below) + # 4. Adjust assertions based on your expected outcomes + # 5. Add custom validation logic for your specific use case + # ============================================================================ + + #Please list the relevant test points. + """ + Test multi-modal input processing and text/audio output generation via OpenAI API. + Deploy Setting: default yaml + Input Modal: text + audio + video + image + Output Modal: text + audio + Input Setting: stream=True + Datasets: single request + """ + # SECTION 1: TEST SETUP AND INITIALIZATION + # ========================================= + # INSTRUCTIONS: Initialize test variables and prepare test environment + # MODIFY: Add any additional test setup required for your scenario + e2e_list = list() + # SECTION 2: TEST DATA GENERATION + # ================================ + # INSTRUCTIONS: Generate or load test data for each input modality + # MODIFY: Replace synthetic generators with your actual data sources + # VIDEO DATA - Generate synthetic video for testing + # FORMAT: data:video/mp4;base64,{base64_encoded_video} + # PARAMETERS: width, height, duration_frames + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + # IMAGE DATA - Generate synthetic image for testing + # FORMAT: data:image/jpeg;base64,{base64_encoded_image} + # PARAMETERS: width, height + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + # AUDIO DATA - Generate synthetic audio for testing + # FORMAT: data:audio/wav;base64,{base64_encoded_audio} + # PARAMETERS: duration_seconds, channels + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" + + # SECTION 3: MESSAGE CONSTRUCTION + # ================================ + # INSTRUCTIONS: Assemble the complete message payload for API request + # MODIFY: Add/remove modalities or change prompt structure as needed + + # USAGE: Construct a message containing all input modalities + # IMPORTANT: Ensure the message structure matches OpenAI API expectations + # CUSTOMIZATION POINTS: + # - system_prompt: Controls the assistant's behavior + # - content_text: The user's text prompt/question + # - *_data_url: URLs for media content (video/image/audio) + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + video_data_url=video_data_url, + image_data_url=image_data_url, + audio_data_url=audio_data_url, + content_text=get_prompt("mix"), + ) + + # SECTION 4: API REQUEST EXECUTION + # ================================= + # INSTRUCTIONS: Make the API call and measure performance + # MODIFY: Add timeout, retry logic, or additional parameters + start_time = time.perf_counter() + chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True) + + #Call using your preferred method and obtain the final audio and text outputs. + ... + + # SECTION 5: OUTPUT VALIDATION + # ============================= + # INSTRUCTIONS: Verify that outputs meet expected criteria + # MODIFY: Adjust validation logic for your specific requirements + + # ASSERTION 1: E2E Validation + # PURPOSE: Verify that the E2E latency is less than the baseline. + current_e2e = time.perf_counter() - start_time + print(f"the request e2e is: {current_e2e}") + e2e_list.append(current_e2e) + + print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}") + + + + # ASSERTION 2: Text Output Validation + # PURPOSE: Verify that text output was generated with keyword content + assert text_content is not None and len(text_content) >= 2, "No text output is generated" + assert any( + keyword in text_content.lower() for keyword in ["square", "quadrate", "sphere", "globe", "circle", "round"] + ), "The output does not contain any of the keywords." + + + # ASSERTION 3: Cross-Modal Consistency + # PURPOSE: Verify text and audio outputs convey the same information + # CUSTOMIZATION: Adjust similarity threshold (0.9) based on accuracy requirements + assert audio_data is not None, "No audio output is generated" + audio_content = merge_base64_and_convert_to_text(audio_data) + print(f"text content is: {text_content}") + print(f"audio content is: {audio_content}") + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" ``` + #### E2E - Offline inference ```python # SPDX-License-Identifier: Apache-2.0 @@ -271,10 +465,3 @@ def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None: has_audio = any(o.final_output_type == "audio" for o in outputs) assert has_audio ``` - -## Checklist before submitting your test files - -1. The file is saved in an appropriate place and the file name is clear. -2. The coding style follows the requirements outlined above. -3. **All test functions have appropriate pytest markers** -4. For tests that need run in CI, please ensure the test is configured under the `./buildkite/` folder. diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 529cab752..10755c6c6 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -763,6 +763,37 @@ omni = Omni(model="your-model", ) ``` +### CPU Offload + +See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md) + +vLLM-Omni provides two offloading strategies to reduce GPU memory usage: + +1. **Model-level offload**: Mutual exclusion between DiT and encoders (only one on GPU at a time) +2. **Layerwise (Blockwise) offload**: Keeps only a single transformer block on GPU at a time with compute-memory overlap + +**Usage:** Enable offload when initializing: +```python +# Model-level offload +omni = Omni(model="your-model", enable_cpu_offload=True) + +# Layerwise offload +omni = Omni(model="your-model", enable_layerwise_offload=True) +``` + +**To support layerwise offloading:** Define the blocks attribute name in your transformer: + +```python +class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attr = "blocks" # Attribute name containing transformer blocks + + def __init__(self): + self.blocks = nn.ModuleList([...]) # Transformer blocks +``` + +**Note:** Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. + + --- ## Troubleshooting diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md new file mode 100644 index 000000000..be46b5b14 --- /dev/null +++ b/docs/features/custom_pipeline.md @@ -0,0 +1,151 @@ +# Custom Pipeline Extension Guide + +Transformer already support Custom Pipeline via +https://github.com/huggingface/diffusers/blob/main/docs/source/en/using-diffusers/custom_pipeline_overview.md + +This guide demonstrates how to use the newly added features for extending vLLM-Omni's diffusion pipeline with custom functionality. + +## Overview + +Three main features enable custom pipeline extension: + +1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality +2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines +3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations + +## Features + +### WorkerWrapperBase + +`WorkerWrapperBase` is a wrapper class that creates `DiffusionWorker` instances with optional extension support. It enables dynamic inheritance, allowing you to add custom methods and functionality to workers without modifying the base worker class. + +**Key capabilities:** +- Dynamic worker class extension via `worker_extension_cls` +- Support for custom pipeline initialization via `custom_pipeline_args` +- Method delegation to underlying worker +- Attribute access forwarding + +**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py` + +### load_format Parameter + +The `load_format` parameter controls how diffusion models are loaded. It supports the following values: + +- **`"default"`**: Standard model loading using the model registry (default behavior) +- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name` +- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately) + +**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py` + +### CustomPipelineWorkerExtension + +`CustomPipelineWorkerExtension` is a mixin class that extends `DiffusionWorker` with the ability to re-initialize the pipeline with a custom implementation. + +**Key method:** +- `re_init_pipeline(custom_pipeline_args)`: Re-initializes the pipeline with custom arguments, properly cleaning up the old pipeline first + +**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py` + +## Usage Example + +### Step 1: Create a Custom Pipeline + +Create a custom pipeline class that extends an existing pipeline. In this example, we extend `QwenImageEditPipeline` to add trajectory tracking: + +```python +# custom_pipeline.py +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline +import torch + +class CustomPipeline(QwenImageEditPipeline): + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__(od_config=od_config, prefix=prefix) + + def forward(self, req, prompt=None, negative_prompt=None, **kwargs): + # Call parent's forward to get normal output + output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs) + + # Add custom trajectory data + actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50) + output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32) + output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32) + + return output +``` + +### Step 2: Use the Custom Pipeline with Omni + +Initialize the `Omni` engine with custom pipeline configuration: + +```python +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# Initialize with custom pipeline +omni = Omni( + model="Qwen/Qwen-Image-Edit", + diffusion_load_format="dummy", # Skip initial loading + custom_pipeline_args={ + "pipeline_class": "custom_pipeline.CustomPipeline" + }, +) + +# Generate with the custom pipeline +outputs = omni.generate( + ... +) + +# Access custom trajectory data +output = outputs[0].request_output[0] +print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}") +print(f"Trajectory latents shape: {output.latents.shape}") +``` + +### Step 3: Run the Example + +The example provided in this directory demonstrates the complete workflow: + +```bash +cd examples/offline_inference/custom_pipeline/image_to_image + +# Run with custom pipeline +python image_edit.py \ + --model Qwen/Qwen-Image-Edit-2511 \ + --image cherry_blossom.jpg \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars" \ + --output output_image_edit.png \ + --num_inference_steps 10 +``` + +## Advanced Usage + +### Custom Worker Extension + +You can create custom worker extensions to add new methods beyond pipeline re-initialization: + +```python +from typing import Any +from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension + +class MyCustomExtension(CustomPipelineWorkerExtension): + def custom_method(self): + """Your custom worker method.""" + return "custom_result" + + def another_method(self, data: Any): + """Another custom method.""" + # Access worker internals via self + return self.model_runner.some_operation(data) + +omni = Omni( + model="Qwen/Qwen-Image-Edit", + diffusion_load_format="dummy", + custom_pipeline_args={ + "pipeline_class": "custom_pipeline.CustomPipeline" + }, + worker_extension_cls=MyCustomExtension, + # Note: worker_extension_cls is an internal parameter + # CustomPipelineWorkerExtension will automatically init pipeline when custom_pipeline_args is provided +) +``` diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 1f82fd608..8786ae964 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -1,101 +1,143 @@ -# CPU Offloading for Diffusion Model +# CPU Offloading for Diffusion Models ## Overview -vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models, allowing you to run larger models on GPUs with limited VRAM: +vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models: -1. **Model-level (Component) Offloading**: Swaps entire model components (DiT transformer, VAE, encoders) between GPU and CPU. -2. **Layerwise (Blockwise) Offloading**: Keeps only a single or a few transformer blocks on GPU at a time, with compute - memory copy overlap. +1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time. +2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap. -Both approaches use pinned memory for faster CPU-GPU transfers. For now, the two offloading strategies could not be used at the same time. +Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority. -## Model-level CPU Offloading +## Model-level (Sequential) Offloading -### Implementation +### How It Works -CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU. +Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks: -**Execution Flow**: -1. Text encoders run on GPU while the DiT transformer is offloaded to CPU. -2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed. -3. After the diffusion step, the transformer returns to CPU and the process repeats as needed. +- **When encoders run**: DiT transformer is offloaded to CPU +- **When DiT runs**: Encoders are offloaded to CPU +- **VAE**: Stays resident on GPU -Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory. +Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed. -### Configuration -You can enable CPU offload in two ways: - -1. **Python API**: set `enable_cpu_offload=True`. +### Usage +**Python API:** ```python from vllm_omni import Omni -if __name__ == "__main__": - - m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True) +m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) ``` -2. **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint. +**CLI:** +```bash +vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload +``` ### Limitations -- Cold start latency increases for over one minute for some models(e.g., Qwen-Image) +- Cold start latency increases +- Adds overhead from CPU-GPU transfers between encoder and denoising phases +- Support single GPU only for now ## Layerwise (Blockwise) Offloading -### Implementation -Layerwise offload operates at transformer block granularity, keeping a single transformer block, or a specified number of blocks, on GPU while others stay in CPU memory. +### How It Works -Unlike full model-wise CPU offload which swaps entire components like DiT and encoders, layerwise offloading applies a sliding window way of loading and offloading weights between gpu and cpu: while block `i` computes, block `i+1` get prefetched asynchronously via pinned memory. In this way, only partial blocks(s) reside on GPU at any moment during inference, so that greatly decrease the memory occupancy. +Layerwise offloading keeps only one transformer block on GPU at a time. -**Execution Flow**: +As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU. -1. During model initialization, all components are loaded to CPU first. Then components other than DiT model(s) in the pipeline, such as VAE and encoders, are moved to GPU. The weights of target transformer blocks are collected as contiguous tensors per layer on CPU with pinned memory; and non-block modules (embeddings, norms, etc) in the DiT model are moved to and stay on GPU. -2. The first block(s) are transferred to GPU during initialization of `LayerwiseOffloader`, before the first denoising step of the very first request. -3. As each block executes, the next block prefetches on a separate CUDA stream for compute - memory copy overlap. After execution, the current block is immediately freed from GPU memory. -4. When the last block completes, the first block prefetches for the next denoising step. +**Execution Flow:** +| Block | Pre-forward Hook | Forward | Post-forward Hook | +|-------|------------------|---------|-------------------| +| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 | +| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 | +| ... | ... | ... | ... | +| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) | -Example of hook executions of a DiT model with n layers, by default keep a single layer on GPU: -| Layer (block) idx | forward pre-hook | forward | forward post-hook | -|-------------------|--------------------------------|------------------|---------------------------| -| layer-0 | prefetch layer 1 (copy stream) | compute layer 0 | free layer-0 gpu weights | -| layer-1 | prefetch layer 2 (copy stream) | compute layer 1 | free layer-1 gpu weights | -| layer-2 | prefetch layer 3 (copy stream) | compute layer 2 | free layer-2 gpu weights | -| ... | ... | ... | ... | -| layer-(n-1) | **prefetch layer 0 (copy stream)** | compute layer (n-1) | free layer (n-1) gpu weights | +Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward. +Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines. -### Configuration - -1. **Python API**: set `enable_layerwise_offload=True` and optionally `layerwise_num_gpu_layers`. +### Usage +**Python API:** ```python from vllm_omni import Omni -if __name__ == "__main__": - m = Omni( - model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", - enable_layerwise_offload=True, - ... - ) +# Text-to-video +m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True) + +# Or image-to-video +m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True) ``` -2. **CLI**: pass `--enable-layerwise-offload` and `--layerwise-num-gpu-layers` to the diffusion service entrypoint. +**CLI:** +```bash +# Text-to-video +vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload -### Supported Models +# Or image-to-video +vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload +``` -| Architecture | Models | Example HF Models | DiT Model Cls | Blocks Attr Name | -|--------------|--------|-------------------|----------|----------| -| `QwenImagePipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | "transformer_blocks" | -| `Wan22Pipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | "blocks" | +### To Support a Model -NOTE: Models must define `_layerwise_offload_blocks_attr` class attribute so that the layerwise offloader finds the target transformer blocks. +Models must define the blocks attribute name for layerwise offloading: + +```python +class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attr = "blocks" # Attribute name containing transformer blocks + + def __init__(self): + self.blocks = nn.ModuleList([...]) # Transformer blocks +``` ### Limitations - Cold start latency increases because of - 1) components are loaded to CPU first at the very first during initialization, + 1) components are loaded to CPU first at the very first during initialization, 2) weight consolidation and pinning -- Performance depends on CPU <-> GPU interconnection (e.g., PCIe bandwidth). +- Performance depends on compute cost and H2D bandwidth as well - Support single GPU only for now + + +### Implementation Notes + +**Module Discovery** + +The offloader automatically discovers pipeline components: + +- **DiT modules**: `transformer`, `transformer_2`, `dit` +- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` +- **VAE**: `vae` + +**Hook System** + +Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code. + +**Backend Architecture** + +``` +OffloadBackend (base class) +├── ModelLevelOffloadBackend → uses SequentialOffloadHook +└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook +``` + +Factory function `get_offload_backend()` selects the appropriate backend based on configuration. + + +## Supported Models + +| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) | +|--------------|----------------|-----------|---------------------|-------------------|-------------| +| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | +| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | +| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` | + +**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_attr` pointing to transformer blocks diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 118415ea8..f8d00e2df 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -30,7 +30,7 @@ The following table shows which models are currently supported by parallelism me | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 0cbe7bb39..bf9968ffa 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -46,6 +46,12 @@ def parse_args(): parser.add_argument("--stage-configs-path", type=str, default=None) parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") + parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)") + parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)") + parser.add_argument( + "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)" + ) + args = parser.parse_args() return args @@ -102,6 +108,10 @@ def main(): seed=52, need_kv_receive=False, num_inference_steps=args.steps, + extra_args={ + "cfg_text_scale": args.cfg_text_scale, + "cfg_img_scale": args.cfg_img_scale, + }, ), ) @@ -158,7 +168,12 @@ def main(): if args.modality == "text2img": params_list[0].max_tokens = 1 # type: ignore # The first stage is a SamplingParam (vllm) if len(params_list) > 1: - params_list[1].num_inference_steps = args.steps # type: ignore # The second stage is an OmniDiffusionSamplingParam + diffusion_params = params_list[1] + diffusion_params.num_inference_steps = args.steps # type: ignore + diffusion_params.extra_args = { # type: ignore + "cfg_text_scale": args.cfg_text_scale, + "cfg_img_scale": args.cfg_img_scale, + } omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) diff --git a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py new file mode 100644 index 000000000..b1c74e2b3 --- /dev/null +++ b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +from typing import Any + +import PIL.Image +import torch + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest + +logger = logging.getLogger(__name__) + + +class CustomPipeline(QwenImageEditPipeline): + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__(od_config=od_config, prefix=prefix) + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + image: PIL.Image.Image | torch.Tensor | None = None, + true_cfg_scale: float = 4.0, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float = 1.0, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + ) -> DiffusionOutput: + """Forward pass for image editing with dummy trajectory data.""" + # Call parent's forward to get the normal output + output = super().forward( + req=req, + prompt=prompt, + negative_prompt=negative_prompt, + image=image, + true_cfg_scale=true_cfg_scale, + height=height, + width=width, + num_inference_steps=num_inference_steps, + sigmas=sigmas, + guidance_scale=guidance_scale, + num_images_per_prompt=num_images_per_prompt, + generator=generator, + latents=latents, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_embeds_mask=negative_prompt_embeds_mask, + output_type=output_type, + attention_kwargs=attention_kwargs, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + max_sequence_length=max_sequence_length, + ) + + # Get actual num_inference_steps used + actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps + + # Create dummy trajectory data + dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32) + + # Inject dummy trajectory data into output + output.trajectory_latents = dummy_trajectory_latents + + return output diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py new file mode 100644 index 000000000..8ad5cbf9b --- /dev/null +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +====================================================== + Qwen Image Editor (via vLLM-Omni) +====================================================== + +Example CLI for Qwen-Image-Edit and compatible models. + +This script edits, combines, or layers images according to a text prompt +using the vLLM Omni diffusion backend. + +Examples +--------- +Single image edit: + python image_edit.py \ + --image input.png \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles saying 'Be Kind'" \ + --output output.png + +Multiple image composition: + python image_edit.py \ + --image input1.png input2.png \ + --prompt "Combine these into a single magical composition" \ + --output combined.png + +Accelerated with cache_dit: + python image_edit.py \ + --image input.png \ + --prompt "Edit description" \ + --cache_backend cache_dit \ + --cache_dit_enable_taylorseer + +Layered RGBA output: + python image_edit.py \ + --model "Qwen/Qwen-Image-Layered" \ + --image input.png \ + --prompt "" \ + --output layered \ + --layers 4 \ + --color-format RGBA +""" + +import argparse +import asyncio +import os +import time +from pathlib import Path + +import torch +from PIL import Image + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +# =========================== +# Argument Parser +# =========================== +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Edit or generate images using Qwen-Image-Edit.") + parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.") + parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).") + parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.") + parser.add_argument("--negative_prompt", type=str, default=None) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--cfg_scale", type=float, default=4.0) + parser.add_argument("--guidance_scale", type=float, default=1.0) + parser.add_argument("--output", type=str, default="output.png") + parser.add_argument("--num_outputs_per_prompt", type=int, default=1) + parser.add_argument("--num_inference_steps", type=int, default=50) + parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"]) + parser.add_argument("--ulysses_degree", type=int, default=1) + parser.add_argument("--ring_degree", type=int, default=1) + parser.add_argument("--tensor_parallel_size", type=int, default=1) + parser.add_argument("--layers", type=int, default=4) + parser.add_argument("--resolution", type=int, default=640) + parser.add_argument("--color-format", type=str, default="RGB") + + # Acceleration + Optimization Options + parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1) + parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0) + parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4) + parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24) + parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3) + parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False) + parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1) + parser.add_argument( + "--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"] + ) + parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"]) + parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2) + parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2]) + parser.add_argument("--enforce_eager", action="store_true") + parser.add_argument("--vae_use_slicing", action="store_true") + parser.add_argument("--vae_use_tiling", action="store_true") + parser.add_argument("--enable-cpu-offload", action="store_true") + + return parser.parse_args() + + +# =========================== +# Core Logic +# =========================== +async def main(): + args = parse_args() + + # ---- Load input images ---- + input_images: list[Image.Image] = [] + for image_path in args.image: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Input image not found: {image_path}") + input_images.append(Image.open(image_path).convert(args.color_format)) + + # Single or multi-image input + input_image: Image.Image | list[Image.Image] = input_images[0] if len(input_images) == 1 else input_images + + # ---- Torch setup ---- + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + parallel_config = DiffusionParallelConfig( + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + ) + + # ---- Cache Config ---- + if args.cache_backend == "cache_dit": + cache_config = { + "Fn_compute_blocks": args.cache_dit_fn_compute_blocks, + "Bn_compute_blocks": args.cache_dit_bn_compute_blocks, + "max_warmup_steps": args.cache_dit_max_warmup_steps, + "residual_diff_threshold": args.cache_dit_residual_diff_threshold, + "max_continuous_cached_steps": args.cache_dit_max_continuous_cached_steps, + "enable_taylorseer": args.cache_dit_enable_taylorseer, + "taylorseer_order": args.cache_dit_taylorseer_order, + "scm_steps_mask_policy": args.cache_dit_scm_steps_mask_policy, + "scm_steps_policy": args.cache_dit_scm_steps_policy, + } + elif args.cache_backend == "tea_cache": + cache_config = {"rel_l1_thresh": args.tea_cache_rel_l1_thresh} + else: + cache_config = None + + # ---- Initialize Omni ---- + omni = Omni( + model=args.model, + vae_use_slicing=args.vae_use_slicing, + vae_use_tiling=args.vae_use_tiling, + cache_backend=args.cache_backend, + cache_config=cache_config, + parallel_config=parallel_config, + enforce_eager=args.enforce_eager, + enable_cpu_offload=args.enable_cpu_offload, + diffusion_load_format="dummy", + custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"}, + ) + + print(">>> Pipeline loaded successfully") + + # ---- Profiling + Info ---- + profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + print(f"\n{'=' * 60}") + print("Generation Configuration") + print(f"Model: {args.model}") + print(f"Inference steps: {args.num_inference_steps}") + print(f"Cache: {args.cache_backend or 'None'}") + print(f"Input(s): {len(input_images)} image(s)") + for idx, img in enumerate(input_images): + print(f" • Image {idx + 1}: size={img.size}") + print( + f"Parallel config: Ulysses={args.ulysses_degree}, Ring={args.ring_degree}, " + f"CFG Parallel={args.cfg_parallel_size}, Tensor Parallel={args.tensor_parallel_size}" + ) + print(f"{'=' * 60}\n") + + if profiler_enabled: + omni.start_profile() + + # ---- Generation ---- + t0 = time.perf_counter() + outputs = omni.generate( + { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "multi_modal_data": {"image": input_image}, + }, + OmniDiffusionSamplingParams( + generator=generator, + true_cfg_scale=args.cfg_scale, + guidance_scale=args.guidance_scale, + num_inference_steps=args.num_inference_steps, + num_outputs_per_prompt=args.num_outputs_per_prompt, + layers=args.layers, + resolution=args.resolution, + ), + ) + t1 = time.perf_counter() + print(f"Generation completed in {t1 - t0:.2f}s") + + if profiler_enabled: + omni.stop_profile() + + # ---- Output Validation ---- + if not outputs: + raise ValueError("No output produced from omni.generate()") + + first_out = outputs[0].request_output[0] + req_out: OmniRequestOutput = first_out + + # Verify trajectory data (from custom pipeline) + print(f"\n{'=' * 60}") + print("TRAJECTORY VERIFICATION:") + assert hasattr(req_out, "latents") and req_out.latents is not None + print(f" ✓ trajectory_latents shape: {req_out.latents.shape}") + print(f" ✓ trajectory_latents dtype: {req_out.latents.dtype}") + print(f" ✓ trajectory_latents mean: {req_out.latents.mean().item():.6f}") + print(f" ✓ trajectory_latents std: {req_out.latents.std().item():.6f}") + assert req_out.latents.shape[0] == args.num_inference_steps, ( + f"Expected {args.num_inference_steps} latent snapshots, got {req_out.latents.shape[0]}" + ) + print(f"{'=' * 60}\n") + + # ---- Save images ---- + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + suffix = output_path.suffix or ".png" + stem = output_path.stem or "output" + + images = req_out.images + if args.num_outputs_per_prompt <= 1: + img = images[0] + if isinstance(img, list): # Layered + for i, sub_img in enumerate(img): + save_path = output_path.parent / f"{stem}_{i}{suffix}" + sub_img.save(save_path) + print(f"Saved layer {i}: {save_path.resolve()}") + else: + img.save(output_path) + print(f"Saved: {output_path.resolve()}") + else: + for idx, img in enumerate(images): + if isinstance(img, list): + for sub_idx, sub_img in enumerate(img): + save_path = output_path.parent / f"{stem}_{idx}_{sub_idx}{suffix}" + sub_img.save(save_path) + print(f"Saved composite: {save_path.resolve()}") + else: + save_path = output_path.parent / f"{stem}_{idx}{suffix}" + img.save(save_path) + print(f"Saved: {save_path.resolve()}") + + +# =========================== +# Entrypoint +# =========================== +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/offline_inference/custom_pipeline/image_to_image/run.sh b/examples/offline_inference/custom_pipeline/image_to_image/run.sh new file mode 100644 index 000000000..1d2706e25 --- /dev/null +++ b/examples/offline_inference/custom_pipeline/image_to_image/run.sh @@ -0,0 +1,6 @@ +python image_edit.py \ + --model Qwen/Qwen-Image-Edit-2511 \ + --image cherry_blossom.jpg \ + --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \ + --output output_image_edit.png \ + --num_inference_steps 10 diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 6ff7b561b..14cff2726 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -300,12 +300,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) return parser.parse_args() @@ -362,7 +356,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index e57f54819..5a91c4c73 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -80,12 +80,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) parser.add_argument( "--cfg-parallel-size", type=int, @@ -139,7 +133,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 5da0c1879..db87abf00 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -113,12 +113,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) parser.add_argument( "--quantization", type=str, @@ -222,7 +216,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index f645a47c0..40fafa100 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -84,12 +84,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) parser.add_argument( "--ulysses-degree", type=int, @@ -155,7 +149,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, diff --git a/requirements/common.txt b/requirements/common.txt index 2331e0f5d..3954b4587 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -13,3 +13,4 @@ openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 sox>=1.5.0 prettytable>=3.8.0 +aenum==3.1.16 diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py new file mode 100644 index 000000000..ea743eadd --- /dev/null +++ b/tests/benchmarks/patch/test_patch.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for patch.py +""" + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest +from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput + +from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput, async_request_openai_chat_omni_completions + +pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu] + + +class MockResponse: + """Mock aiohttp response for testing""" + + def __init__(self, status, chunks, delay_between_chunks=0): + self.status = status + self.reason = "OK" if status == 200 else "Error" + self._chunks = chunks + self._delay = delay_between_chunks + self.content = self + + async def iter_any(self): + for chunk in self._chunks: + if self._delay > 0: + await asyncio.sleep(self._delay) + yield chunk + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + +def create_sse_chunk(data_dict): + """Helper to create SSE formatted chunk""" + return f"data: {json.dumps(data_dict)}\n\n".encode() + + +# ============================================================================ +# output_tokens Tests +# ============================================================================ + + +@pytest.mark.asyncio +async def test_output_tokens_assigned_with_metrics(): + """Test that output.output_tokens is assigned when metrics are present""" + # Arrange + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with metrics + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + "metrics": {"num_tokens_out": 42, "num_tokens_in": 10}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert output.output_tokens == 42, "output_tokens should be assigned from metrics" + assert output.generated_text == "Hello world" + + +@pytest.mark.asyncio +async def test_output_tokens_not_assigned_without_metrics(): + """Test that output.output_tokens defaults to 0 when no metrics present""" + # Arrange + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks without metrics + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert output.output_tokens == 0, "output_tokens should default to 0 when no metrics" + assert output.generated_text == "Hello world" + + +@pytest.mark.asyncio +async def test_output_tokens_assigned_multiple_metrics(): + """Test that output.output_tokens is updated with the latest metrics value""" + # Arrange + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with multiple metrics updates + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + "metrics": {"num_tokens_out": 5}, + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + "metrics": {"num_tokens_out": 10}, + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": "!"}}], + "modality": "text", + "metrics": {"num_tokens_out": 15}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert output.output_tokens == 15, "output_tokens should be updated to the latest value" + assert output.generated_text == "Hello world!" + + +@pytest.mark.asyncio +async def test_output_tokens_with_audio_and_text(): + """Test output_tokens assignment in mixed audio and text response""" + # Arrange + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with both audio and text, with metrics + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Text response"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": ""}}], + "modality": "audio", + } + ), + create_sse_chunk( + { + "modality": "text", + "metrics": {"num_tokens_out": 25, "num_tokens_in": 10}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert output.output_tokens == 25, "output_tokens should be assigned even with audio modality" + + +@pytest.mark.asyncio +async def test_output_tokens_with_missing_num_tokens_out(): + """Test that output_tokens defaults to 0 when num_tokens_out is missing in metrics""" + # Arrange + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with metrics but without num_tokens_out + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + "metrics": {"num_tokens_in": 10}, # Missing num_tokens_out + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert output.output_tokens == 0, "output_tokens should default to 0 when num_tokens_out is missing" + + +@pytest.mark.asyncio +async def test_output_tokens_initialization(): + """Test that MixRequestFuncOutput initializes output_tokens correctly""" + # Arrange & Act + output = MixRequestFuncOutput() + + # Assert + assert hasattr(output, "output_tokens"), "MixRequestFuncOutput should have output_tokens attribute" + assert output.output_tokens == 0, "output_tokens should be initialized to 0" + + +# ============================================================================ +# text_latency Tests +# ============================================================================ + + +class TestTextLatencyAttribute: + """Tests for text_latency attribute existence and assignment""" + + def test_mix_request_func_output_has_text_latency(self): + """Test that MixRequestFuncOutput has text_latency attribute""" + output = MixRequestFuncOutput() + assert hasattr(output, "text_latency"), "MixRequestFuncOutput should have text_latency attribute" + + def test_text_latency_initial_value(self): + """Test that text_latency initializes to a default value""" + output = MixRequestFuncOutput() + # Check if attribute exists and has a value (should be 0.0 or similar default) + text_latency = getattr(output, "text_latency", None) + assert text_latency is not None or hasattr(output, "text_latency"), "text_latency attribute should exist" + + @pytest.mark.asyncio + async def test_text_latency_assigned_with_text_response(self): + """Test that text_latency is assigned when text response is received""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with text modality + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert hasattr(output, "text_latency"), "Output should have text_latency attribute" + assert output.text_latency > 0, "text_latency should be greater than 0 for text response" + + @pytest.mark.asyncio + async def test_text_latency_updated_with_multiple_text_chunks(self): + """Test that text_latency is updated as more text chunks arrive""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with small delays to simulate streaming + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "First"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " second"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " third"}}], + "modality": "text", + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert hasattr(output, "text_latency"), "Output should have text_latency attribute" + assert output.text_latency > 0, "text_latency should accumulate" + assert output.generated_text == "First second third" + + @pytest.mark.asyncio + async def test_text_latency_with_only_audio_response(self): + """Test text_latency behavior when only audio is received""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response chunks with only audio modality (no text) + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": ""}}], + "modality": "audio", + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert hasattr(output, "text_latency"), "Output should have text_latency attribute even with audio-only" + # text_latency should either be 0 or the initial value since no text was processed + assert output.text_latency >= 0, "text_latency should be non-negative" + + @pytest.mark.asyncio + async def test_text_latency_not_affected_by_metrics(self): + """Test that text_latency is independent of metrics data""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response with text and metrics + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Response text"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": ""}}], + "modality": "text", + "metrics": {"num_tokens_out": 100, "num_tokens_in": 20}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert hasattr(output, "text_latency"), "text_latency should exist" + assert output.text_latency > 0, "text_latency should be set for text response" + # Verify metrics are also present + assert output.output_tokens == 100, "metrics should not affect text_latency" + + @pytest.mark.asyncio + async def test_text_latency_mixed_modalities(self): + """Test text_latency with mixed text and audio modalities""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + # Create response with both text and audio + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Text"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": ""}}], + "modality": "audio", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " more text"}}], + "modality": "text", + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert hasattr(output, "text_latency"), "text_latency should exist with mixed modalities" + assert output.text_latency > 0, "text_latency should be set when text is present" + assert output.generated_text == "Text more text" + + @pytest.mark.asyncio + async def test_text_latency_value_consistency(self): + """Test that text_latency matches latency minus ttft relationship""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = AsyncMock() + mock_session.post = MagicMock(return_value=mock_response) + + # Act + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + # Assert + assert output.success is True + assert hasattr(output, "text_latency"), "text_latency should exist" + assert hasattr(output, "ttft"), "ttft should exist" + assert hasattr(output, "latency"), "latency should exist" + # text_latency should be between ttft and total latency + assert output.ttft <= output.text_latency <= output.latency, ( + "text_latency should be between ttft and total latency" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 6e57355a2..73cfa104f 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -36,7 +36,7 @@ def mock_od_config(): @pytest.fixture def mock_gpu_worker(mock_od_config): """Create a DiffusionWorker with mocked initialization.""" - with patch.object(DiffusionWorker, "init_device"): + with patch.object(DiffusionWorker, "init_device"), patch.object(DiffusionWorker, "load_model"): worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config) # Mock the model_runner with pipeline worker.model_runner = Mock() @@ -81,18 +81,18 @@ class TestDiffusionWorkerLoadWeights: class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker): + def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" - # Setup memory info mocks - # Before sleep: 1GB free - # After sleep: 3GB free (freed 2GB) - mock_platform.get_free_memory.side_effect = [ - 1 * 1024**3, # Before sleep - 3 * 1024**3, # After sleep + # Setup process-scoped memory mocks + # Before sleep: 3GB used + # After sleep: 1GB used (freed 2GB) + mock_get_process_memory.side_effect = [ + 3 * 1024**3, + 1 * 1024**3, ] - mock_platform.get_device_total_memory.return_value = 8 * 1024**3 # Setup allocator mock mock_allocator = Mock() @@ -108,16 +108,16 @@ class TestDiffusionWorkerSleep: # Verify buffers were NOT saved (level 1 doesn't save buffers) assert len(mock_gpu_worker._sleep_saved_buffers) == 0 + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker): + def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" - # Setup memory info mocks - mock_platform.get_free_memory.side_effect = [ - 1 * 1024**3, # Before sleep - 5 * 1024**3, # After sleep (freed 4GB) + # Setup process-scoped memory mocks + mock_get_process_memory.side_effect = [ + 5 * 1024**3, # Before sleep + 1 * 1024**3, # After sleep (freed 4GB) ] - mock_platform.get_device_total_memory.return_value = 8 * 1024**3 # Setup allocator mock mock_allocator = Mock() @@ -146,16 +146,22 @@ class TestDiffusionWorkerSleep: assert "buffer1" in mock_gpu_worker._sleep_saved_buffers assert "buffer2" in mock_gpu_worker._sleep_saved_buffers + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") @patch("vllm.device_allocator.cumem.CuMemAllocator") - def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker): + def test_sleep_memory_freed_validation( + self, + mock_allocator_class, + mock_platform, + mock_get_process_memory, + mock_gpu_worker, + ): """Test that sleep validates memory was actually freed.""" - # Simulate memory increase (should trigger assertion error) - mock_platform.get_free_memory.side_effect = [ - 3 * 1024**3, # Before sleep: 3GB free - 1 * 1024**3, # After sleep: 1GB free (negative freed!) + # Simulate process memory increase (should trigger assertion error) + mock_get_process_memory.side_effect = [ + 1 * 1024**3, # Before sleep: 1GB used + 3 * 1024**3, # After sleep: 3GB used (negative freed) ] - mock_platform.get_device_total_memory.return_value = 8 * 1024**3 mock_allocator = Mock() mock_allocator_class.get_instance = Mock(return_value=mock_allocator) @@ -165,6 +171,33 @@ class TestDiffusionWorkerSleep: with pytest.raises(AssertionError, match="Memory usage increased after sleeping"): mock_gpu_worker.sleep(level=1) + @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") + @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + @patch("vllm.device_allocator.cumem.CuMemAllocator") + def test_sleep_falls_back_to_device_memory_when_nvml_unavailable( + self, + mock_allocator_class, + mock_platform, + mock_get_process_memory, + mock_gpu_worker, + ): + """Test sleep uses device-scoped fallback when NVML is unavailable.""" + mock_get_process_memory.side_effect = [None, None] + mock_platform.get_free_memory.side_effect = [ + 1 * 1024**3, # Before sleep + 3 * 1024**3, # After sleep + ] + mock_platform.get_device_total_memory.return_value = 8 * 1024**3 + + mock_allocator = Mock() + mock_allocator_class.get_instance = Mock(return_value=mock_allocator) + mock_allocator.sleep = Mock() + + result = mock_gpu_worker.sleep(level=1) + + mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) + assert result is True + class TestDiffusionWorkerWakeUp: """Test DiffusionWorker.wake_up method.""" diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py new file mode 100644 index 000000000..2023a89e1 --- /dev/null +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -0,0 +1,540 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for WorkerWrapperBase class. + +This module tests the WorkerWrapperBase implementation: +- Initialization with and without worker extensions +- Custom pipeline initialization +- Method delegation via execute_method +- Attribute delegation via __getattr__ +- Dynamic worker class extension +""" + +from typing import Any +from unittest.mock import Mock, patch + +import pytest + +from vllm_omni.diffusion.worker.diffusion_worker import ( + CustomPipelineWorkerExtension, + DiffusionWorker, + WorkerWrapperBase, +) + +# ------------------------------------------------------------------------- +# Fixtures +# ------------------------------------------------------------------------- + + +@pytest.fixture +def mock_od_config(): + """Create a mock OmniDiffusionConfig for use in tests.""" + config = Mock() + config.num_gpus = 1 + config.master_port = 12345 + config.enable_sleep_mode = False + config.cache_backend = None + config.cache_config = None + config.model = "test-model" + config.diffusion_load_format = None + config.dtype = "float32" + config.max_cpu_loras = 0 + config.lora_path = None + config.lora_scale = 1.0 + return config + + +class TestExtension: + """Simple test extension adding one custom method.""" + + def custom_method(self): + return "extension_method" + + +class MockCustomPipeline: + """Mock custom pipeline for testing.""" + + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + def __call__(self, *args, **kwargs): + return "pipeline_output" + + +# ------------------------------------------------------------------------- +# Tests: Initialization +# ------------------------------------------------------------------------- + + +class TestWorkerWrapperBaseInitialization: + """Test WorkerWrapperBase initialization behavior.""" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_basic_initialization(self, mock_worker_init, mock_od_config): + """Test basic initialization without extensions.""" + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + ) + + assert wrapper.gpu_id == 0 + assert wrapper.od_config == mock_od_config + assert wrapper.base_worker_class == DiffusionWorker + assert wrapper.worker_extension_cls is None + assert wrapper.custom_pipeline_args is None + assert wrapper.worker is not None + + mock_worker_init.assert_called_once_with( + local_rank=0, + rank=0, + od_config=mock_od_config, + ) + + +# ------------------------------------------------------------------------- +# Tests: Worker Extension Functionality +# ------------------------------------------------------------------------- + + +class TestWorkerWrapperBaseExtension: + """Test WorkerWrapperBase worker extension functionality.""" + + def test_prepare_worker_class_without_extension(self, mock_od_config): + """Test _prepare_worker_class without a worker extension.""" + with patch.object(DiffusionWorker, "__init__", return_value=None): + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + ) + worker_class = wrapper._prepare_worker_class() + assert worker_class == DiffusionWorker + + def test_prepare_worker_class_with_extension_class(self, mock_od_config): + """Test _prepare_worker_class with an explicit extension class.""" + + class TestExtension: + def custom_method(self): + return "extension_method" + + with patch.object(DiffusionWorker, "__init__", return_value=None): + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=TestExtension, + ) + + assert hasattr(wrapper.worker.__class__, "custom_method") + assert TestExtension in wrapper.worker.__class__.__bases__ + + @patch("vllm.utils.import_utils.resolve_obj_by_qualname") + def test_prepare_worker_class_with_extension_string(self, mock_resolve, mock_od_config): + """Test _prepare_worker_class with worker extension as string.""" + mock_resolve.return_value = TestExtension + + with patch.object(DiffusionWorker, "__init__", return_value=None): + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension", + ) + + assert hasattr(wrapper.worker.__class__, "custom_method") + + +# ------------------------------------------------------------------------- +# Tests: Method Delegation +# ------------------------------------------------------------------------- + + +class TestWorkerWrapperBaseDelegation: + """Test WorkerWrapperBase delegation to wrapped worker.""" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_generate_delegation(self, mock_worker_init, mock_od_config): + """Test that generate() delegates to worker.generate().""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + mock_output = Mock() + wrapper.worker.generate = Mock(return_value=mock_output) + + mock_requests = [Mock()] + result = wrapper.generate(mock_requests) + + wrapper.worker.generate.assert_called_once_with(mock_requests) + assert result == mock_output + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_execute_model_delegation(self, mock_worker_init, mock_od_config): + """Test that execute_model() delegates to worker.execute_model().""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + mock_output = Mock() + wrapper.worker.execute_model = Mock(return_value=mock_output) + + mock_reqs = [Mock()] + result = wrapper.execute_model(mock_reqs, mock_od_config) + + wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config) + assert result == mock_output + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_load_weights_delegation(self, mock_worker_init, mock_od_config): + """Test that load_weights() delegates to worker.load_weights().""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + expected_result = {"weight1", "weight2"} + wrapper.worker.load_weights = Mock(return_value=expected_result) + + mock_weights = [("weight1", Mock()), ("weight2", Mock())] + result = wrapper.load_weights(mock_weights) + + wrapper.worker.load_weights.assert_called_once_with(mock_weights) + assert result == expected_result + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_sleep_delegation(self, mock_worker_init, mock_od_config): + """Test that sleep() delegates to worker.sleep().""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.sleep = Mock(return_value=True) + result = wrapper.sleep(level=1) + + wrapper.worker.sleep.assert_called_once_with(1) + assert result is True + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_wake_up_delegation(self, mock_worker_init, mock_od_config): + """Test that wake_up() delegates to worker.wake_up().""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.wake_up = Mock(return_value=True) + + result = wrapper.wake_up(tags=["weights"]) + wrapper.worker.wake_up.assert_called_once_with(["weights"]) + assert result is True + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_shutdown_delegation(self, mock_worker_init, mock_od_config): + """Test that shutdown() delegates to worker.shutdown().""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.shutdown = Mock(return_value=None) + + result = wrapper.shutdown() + wrapper.worker.shutdown.assert_called_once() + assert result is None + + +# ------------------------------------------------------------------------- +# Tests: execute_method +# ------------------------------------------------------------------------- + + +class TestWorkerWrapperBaseExecuteMethod: + """Test WorkerWrapperBase.execute_method functionality.""" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_execute_method_success(self, mock_worker_init, mock_od_config): + """Test execute_method successfully calls worker method.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.test_method = Mock(return_value="method_result") + + result = wrapper.execute_method("test_method", "arg1", kwarg1="value1") + + wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1") + assert result == "method_result" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_execute_method_with_no_args(self, mock_worker_init, mock_od_config): + """Test execute_method with no arguments.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.no_args_method = Mock(return_value="no_args_result") + + result = wrapper.execute_method("no_args_method") + wrapper.worker.no_args_method.assert_called_once_with() + assert result == "no_args_result" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_execute_method_error(self, mock_worker_init, mock_od_config): + """Test execute_method raises exception on error.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.error_method = Mock(side_effect=RuntimeError("Test error")) + + with pytest.raises(RuntimeError, match="Test error"): + wrapper.execute_method("error_method") + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_execute_method_invalid_type(self, mock_worker_init, mock_od_config): + """Test execute_method with invalid method type.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + + with pytest.raises(AssertionError, match="Method must be str"): + wrapper.execute_method(b"bytes_method") + + +# ------------------------------------------------------------------------- +# Tests: __getattr__ delegation +# ------------------------------------------------------------------------- + + +class TestWorkerWrapperBaseGetAttr: + """Test WorkerWrapperBase.__getattr__ delegation.""" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_getattr_delegation(self, mock_worker_init, mock_od_config): + """Test __getattr__ delegates to worker attributes.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.custom_attribute = "test_value" + assert wrapper.custom_attribute == "test_value" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_getattr_method_access(self, mock_worker_init, mock_od_config): + """Test __getattr__ delegates to worker methods.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + wrapper.worker.custom_method = Mock(return_value="method_result") + + result = wrapper.custom_method() + wrapper.worker.custom_method.assert_called_once() + assert result == "method_result" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_getattr_missing_attribute(self, mock_worker_init, mock_od_config): + """Test __getattr__ raises AttributeError for missing attributes.""" + wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) + with pytest.raises(AttributeError): + _ = wrapper.nonexistent_attribute + + +# ------------------------------------------------------------------------- +# Tests: Edge Cases +# ------------------------------------------------------------------------- + + +class TestWorkerWrapperBaseEdgeCases: + """Test WorkerWrapperBase edge cases and special scenarios.""" + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_extension_conflict_warning(self, mock_worker_init, mock_od_config, caplog): + """Test a warning is logged when an extension conflicts with worker.""" + + class ConflictExtension: + def load_model(self): + return "extension_load_model" + + with patch.object(DiffusionWorker, "load_model"): + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=ConflictExtension, + ) + assert wrapper.worker is not None + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_multiple_extensions_same_class(self, mock_worker_init, mock_od_config): + """Test that applying same extension twice doesn't duplicate it.""" + + class TestExtension: + def custom_method(self): + return "extension" + + with patch.object(DiffusionWorker, "__init__", return_value=None): + wrapper1 = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=TestExtension, + ) + wrapper2 = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=TestExtension, + ) + + assert hasattr(wrapper1.worker, "custom_method") + assert hasattr(wrapper2.worker, "custom_method") + + +# ------------------------------------------------------------------------- +# Tests: CustomPipelineWorkerExtension +# ------------------------------------------------------------------------- + + +class TestCustomPipelineWorkerExtension: + """Test CustomPipelineWorkerExtension functionality.""" + + @patch("torch.cuda.empty_cache") + @patch("gc.collect") + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_re_init_pipeline_basic(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + """Test basic re_init_pipeline functionality.""" + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=CustomPipelineWorkerExtension, + ) + + # Setup mock model_runner and pipeline + mock_model_runner = Mock() + mock_pipeline = Mock() + mock_model_runner.pipeline = mock_pipeline + wrapper.worker.model_runner = mock_model_runner + wrapper.worker.init_lora_manager = Mock() + wrapper.worker.load_model = Mock() + + custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + + # Call re_init_pipeline + wrapper.worker.re_init_pipeline(custom_args) + + # Verify load_model was called with correct arguments + wrapper.worker.load_model.assert_called_once_with( + load_format="custom_pipeline", + custom_pipeline_name="tests.diffusion.test_worker_wrapper_base.MockCustomPipeline", + ) + wrapper.worker.init_lora_manager.assert_called_once() + + @patch("torch.cuda.empty_cache") + @patch("gc.collect") + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_re_init_pipeline_cleanup(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + """Test that re_init_pipeline properly cleans up old pipeline.""" + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=CustomPipelineWorkerExtension, + ) + + # Setup mock model_runner with pipeline + mock_model_runner = Mock() + mock_pipeline = Mock() + mock_model_runner.pipeline = mock_pipeline + wrapper.worker.model_runner = mock_model_runner + wrapper.worker.init_lora_manager = Mock() + wrapper.worker.load_model = Mock() + + custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + + # Call re_init_pipeline + wrapper.worker.re_init_pipeline(custom_args) + + # Verify cleanup was performed + mock_gc_collect.assert_called_once() + mock_empty_cache.assert_called_once() + + @patch("torch.cuda.empty_cache") + @patch("gc.collect") + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_re_init_pipeline_none_pipeline(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + """Test re_init_pipeline when pipeline is None.""" + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=CustomPipelineWorkerExtension, + ) + + # Setup mock model_runner with None pipeline + mock_model_runner = Mock() + mock_model_runner.pipeline = None + wrapper.worker.model_runner = mock_model_runner + wrapper.worker.init_lora_manager = Mock() + wrapper.worker.load_model = Mock() + + custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + + # Should not raise an error + wrapper.worker.re_init_pipeline(custom_args) + + # Verify load_model was still called + wrapper.worker.load_model.assert_called_once() + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_custom_pipeline_args_initialization(self, mock_worker_init, mock_od_config): + """Test initialization with custom_pipeline_args calls re_init_pipeline.""" + custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + + with patch.object(DiffusionWorker, "__init__", return_value=None): + with patch.object(WorkerWrapperBase, "_prepare_worker_class") as mock_prepare: + # Create a mock worker class with re_init_pipeline + mock_worker_class = Mock() + mock_worker_instance = Mock() + mock_worker_instance.re_init_pipeline = Mock() + mock_worker_class.return_value = mock_worker_instance + mock_prepare.return_value = mock_worker_class + + _ = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + custom_pipeline_args=custom_args, + ) + + # Verify re_init_pipeline was called with custom_pipeline_args + mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args) + + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_custom_pipeline_with_explicit_extension(self, mock_worker_init, mock_od_config): + """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided.""" + + class CustomExtension: + def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]): + return "custom_re_init_pipeline" + + def custom_extension_method(self): + return "custom_extension_method" + + custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + + with patch.object(DiffusionWorker, "__init__", return_value=None): + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=CustomExtension, + custom_pipeline_args=custom_args, + ) + + # Should still have the explicitly provided extension + assert CustomExtension in wrapper.worker.__class__.__bases__ + assert hasattr(wrapper.worker, "custom_extension_method") + + @patch("torch.cuda.empty_cache") + @patch("gc.collect") + @patch.object(DiffusionWorker, "__init__", return_value=None) + def test_re_init_pipeline_multiple_calls(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config): + """Test calling re_init_pipeline multiple times.""" + wrapper = WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + worker_extension_cls=CustomPipelineWorkerExtension, + ) + + # Setup mock model_runner + mock_model_runner = Mock() + mock_pipeline1 = Mock() + mock_pipeline2 = Mock() + mock_model_runner.pipeline = mock_pipeline1 + wrapper.worker.model_runner = mock_model_runner + wrapper.worker.init_lora_manager = Mock() + wrapper.worker.load_model = Mock() + + # First call + custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + wrapper.worker.re_init_pipeline(custom_args1) + + # Update pipeline for second call + mock_model_runner.pipeline = mock_pipeline2 + + # Second call + custom_args2 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} + wrapper.worker.re_init_pipeline(custom_args2) + + # Verify load_model was called twice with different pipelines + assert wrapper.worker.load_model.call_count == 2 + assert wrapper.worker.init_lora_manager.call_count == 2 diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py new file mode 100644 index 000000000..a5ebc2227 --- /dev/null +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from collections import deque +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.v1.request import RequestStatus + +from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) +from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class DummyWaitingQueue(list): + def prepend_requests(self, requests): + self[:0] = list(requests) + + def add_request(self, request): + self.append(request) + + +def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None): + return SimpleNamespace( + request_id=req_id, + external_req_id=external_req_id or req_id, + status=status, + prompt_token_ids=[], + num_computed_tokens=0, + additional_information=None, + ) + + +@pytest.fixture +def build_adapter(monkeypatch): + def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2): + connector = MagicMock() + connector.stage_id = stage_id + connector.get.return_value = None + connector.put.return_value = (True, 1, {}) + + def _fake_base_init(self, config): + self.config = config + self._pending_load_reqs = {} + self._finished_load_reqs = set() + self._pending_save_reqs = {} + self._finished_save_reqs = set() + self.stop_event = threading.Event() + self.lock = threading.Lock() + + monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init) + monkeypatch.setattr( + OmniChunkTransferAdapter, + "create_connector", + classmethod(lambda cls, _model_config: connector), + ) + + model_config = SimpleNamespace(worker_type=model_mode) + scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) + adapter = OmniChunkTransferAdapter( + SimpleNamespace(model_config=model_config, scheduler_config=scheduler_config) + ) + return adapter, connector + + return _build + + +@pytest.mark.parametrize( + ("raw_cfg", "expected_name", "expected_extra"), + [ + (None, "SharedMemoryConnector", {}), + (SimpleNamespace(name="YuanrongConnector", extra={"k": "v"}), "YuanrongConnector", {"k": "v"}), + ], +) +def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, expected_extra): + captured = {} + + def _fake_create(spec): + captured["spec"] = spec + return "ok" + + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter" + ".OmniConnectorFactory.create_connector", + _fake_create, + ) + + model_config = SimpleNamespace(stage_connector_config=raw_cfg) if raw_cfg is not None else SimpleNamespace() + connector = OmniChunkTransferAdapter.create_connector(model_config) + + assert connector == "ok" + assert isinstance(captured["spec"], ConnectorSpec) + assert captured["spec"].name == expected_name + assert captured["spec"].extra == expected_extra + + +def test_load_poll(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1") + + adapter.load_async(request) + payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True} + connector.get.return_value = (payload, 16) + adapter._poll_single_request("req-1") + + connector.get.assert_called_once_with("1", "2", "external-1_1_0") + assert request.additional_information == payload + assert adapter.get_req_chunk["req-1"] == 1 + assert "req-1" in adapter._finished_load_reqs + assert "req-1" in adapter.finished_requests + assert "req-1" not in adapter._pending_load_reqs + + +def test_save_async(build_adapter): + adapter, _ = build_adapter(stage_id=1) + request = SimpleNamespace(external_req_id="external-1") + + adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False} + adapter.save_async(pooling_output=None, request=request) + adapter.custom_process_next_stage_input_func = lambda **kwargs: {} + adapter.save_async(pooling_output=None, request=request) + + assert adapter.put_req_chunk["external-1"] == 1 + queued = adapter._pending_save_reqs["external-1"] + assert len(queued) == 1 + assert queued[0]["put_key"] == "external-1_1_0" + + +def test_update_request_payload(build_adapter): + adapter, _ = build_adapter() + + adapter._update_request_payload("ext", {"h": torch.tensor([[1.0]]), "codes": [1], "finished": False}) + merged = adapter._update_request_payload("ext", {"h": torch.tensor([[2.0]]), "codes": [2], "finished": True}) + + assert torch.equal(merged["h"], torch.tensor([[1.0], [2.0]])) + assert merged["codes"] == [1, 2] + assert merged["finished"] is True + + +def test_process_and_restore_queues(build_adapter): + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + waiting_req = _req("w1", RequestStatus.WAITING) + running_req = _req("r1", RequestStatus.RUNNING) + waiting_queue = DummyWaitingQueue([waiting_req]) + running_queue = [running_req] + + adapter.process_pending_chunks(waiting_queue, running_queue) + assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK + assert running_req.status == RequestStatus.WAITING_FOR_CHUNK + assert waiting_queue == [] + assert running_queue == [] + + adapter.restore_queues(waiting_queue, running_queue) + assert waiting_queue == [waiting_req] + assert running_queue == [running_req] + assert adapter.waiting_for_chunk_waiting_requests == deque() + assert adapter.waiting_for_chunk_running_requests == deque() + + +def test_postprocess_scheduler_output(build_adapter): + adapter, _ = build_adapter() + adapter.requests_with_ready_chunks = {"new-ready", "cached-ready", "leftover"} + + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[SimpleNamespace(req_id="new-ready")], + scheduled_cached_reqs=SimpleNamespace(req_ids=["cached-ready", "missing"]), + ) + requests = {"cached-ready": SimpleNamespace(additional_information={"k": "v"})} + + adapter.postprocess_scheduler_output(scheduler_output, requests) + + cached_info = scheduler_output.scheduled_cached_reqs.additional_information + assert cached_info["cached-ready"] == {"k": "v"} + assert cached_info["missing"] is None + assert adapter.requests_with_ready_chunks == {"leftover"} diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index 87a9e0a9e..1a8b78969 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py @@ -17,13 +17,15 @@ if str(REPO_ROOT) not in sys.path: from vllm_omni import Omni # Models to test and expected saved memory in MB, correspondingly -MODELS_SAVED_MEMORY_MB = {"riverclouds/qwen_image_random": 4500} +MODELS_SAVED_MEMORY_MB = { + "riverclouds/qwen_image_random": 4500, + # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": 45000, +} def run_inference( model_name: str, layerwise_offload: bool = False, - num_gpu_layers: int = 1, num_inference_steps: int = 3, ) -> float: # For now, only support on GPU, so apply torch.cuda operations here @@ -36,7 +38,6 @@ def run_inference( m = Omni( model=model_name, enable_layerwise_offload=layerwise_offload, - layerwise_num_gpu_layers=num_gpu_layers, boundary_ratio=0.875, flow_shift=5.0, ) @@ -83,16 +84,12 @@ def test_layerwise_offload_diffusion_model(model_name: str): cleanup_dist_env_and_memory() # Run with layerwise offloading (1 layer on device) - layerwise_offload_peak_memory = run_inference(model_name, layerwise_offload=True, num_gpu_layers=1) + layerwise_offload_peak_memory = run_inference(model_name, layerwise_offload=True) cleanup_dist_env_and_memory() - - # Run with 2 layers on device - layerwise_offload_two_layers_peak = run_inference(model_name, layerwise_offload=True, num_gpu_layers=2) except Exception: pytest.fail("Inference failed") print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB") - print(f"Layerwise offload peak memory (2 GPU layers): {layerwise_offload_two_layers_peak} MB") print(f"No offload peak memory: {no_offload_peak_memory} MB") # Verify that layerwise offloading significantly reduces memory usage @@ -101,10 +98,3 @@ def test_layerwise_offload_diffusion_model(model_name: str): f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB " f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB" ) - - # Verify that 2 GPU layers uses more memory than 1 GPU layer - # But not excessively more (should be a reasonable increase) - assert layerwise_offload_peak_memory < layerwise_offload_two_layers_peak, ( - f"1 GPU layer peak {layerwise_offload_peak_memory} MB should be < " - f"2 GPU layers peak {layerwise_offload_two_layers_peak} MB" - ) diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py index e2cb09f30..33960f560 100644 --- a/tests/e2e/offline_inference/test_t2i_model.py +++ b/tests/e2e/offline_inference/test_t2i_model.py @@ -34,10 +34,17 @@ elif current_omni_platform.is_rocm(): @pytest.mark.core_model +@pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 2}) @pytest.mark.parametrize("model_name", models) -def test_diffusion_model(model_name: str): +def test_diffusion_model(model_name: str, run_level): + if run_level == "core_model" and model_name != "riverclouds/qwen_image_random": + pytest.skip() + + if run_level == "advanced_model" and model_name == "riverclouds/qwen_image_random": + pytest.skip() + m = None try: m = Omni(model=model_name) diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index 79fdb4699..c43a5d224 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -154,7 +154,7 @@ def _run_zimage_generate( cleanup_dist_env_and_memory() -@pytest.mark.core_model +@pytest.mark.advanced_model @pytest.mark.diffusion @pytest.mark.parallel @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py index 85b80cafb..ac5d81ce0 100644 --- a/tests/e2e/online_serving/test_images_generations_lora.py +++ b/tests/e2e/online_serving/test_images_generations_lora.py @@ -145,7 +145,7 @@ def _basic_payload() -> dict: } -@pytest.mark.core_model +@pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_images_generations_per_request_lora_switching(omni_server: OmniServer, tmp_path: Path) -> None: diff --git a/tests/e2e/stage_configs/qwen3_omni_thinker_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_thinker_ci.yaml index 89d03966d..a6b4404d7 100644 --- a/tests/e2e/stage_configs/qwen3_omni_thinker_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_thinker_ci.yaml @@ -10,7 +10,7 @@ stage_args: worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 - enforce_eager: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 5f36357b8..83671b713 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -184,7 +184,7 @@ def calculate_metrics( total_input += input_requests[i].prompt_len tpot = 0 if output_len > 1: - latency_minus_ttft = outputs[i].latency - outputs[i].ttft + latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft tpot = latency_minus_ttft / (output_len - 1) tpots.append(tpot) # Note: if output_len <= 1, we regard tpot as 0 for goodput diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 6066be44a..5a03a8069 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -72,6 +72,7 @@ class MixRequestFuncOutput(RequestFuncOutput): audio_duration: float = 0.0 audio_frames: int = 0 audio_rtf: float = 0.0 + text_latency: float = 0.0 async def async_request_openai_chat_omni_completions( @@ -149,6 +150,7 @@ async def async_request_openai_chat_omni_completions( output.itl.append(timestamp - most_recent_timestamp) generated_text += content or "" most_recent_timestamp = timestamp + output.text_latency = timestamp - st elif modality == "audio": if output.audio_ttfp == 0.0: output.audio_ttfp = timestamp - st @@ -162,8 +164,8 @@ async def async_request_openai_chat_omni_completions( else: generated_audio = generated_audio + seg - elif usage := data.get("usage"): - output.output_tokens = usage.get("completion_tokens") + if metrics := data.get("metrics"): + output.output_tokens = metrics.get("num_tokens_out", 0) output.latency = timestamp - st output.generated_text = generated_text diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 08e598d99..80dcaf7b6 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -28,6 +28,7 @@ class OmniModelConfig(ModelConfig): (default: "thinker") model_arch: Model architecture name (default: "Qwen2_5OmniForConditionalGeneration") + worker_type: Model Type, e.g., "ar" or "generation" engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. @@ -46,6 +47,7 @@ class OmniModelConfig(ModelConfig): async_chunk: bool = False model_stage: str = "thinker" model_arch: str = "Qwen2_5OmniForConditionalGeneration" + worker_type: str | None = None engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 760bacf00..7b0a1e54d 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,6 +1,5 @@ from __future__ import annotations -import importlib from collections import defaultdict from dataclasses import asdict, dataclass from time import time @@ -20,9 +19,9 @@ from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.output import OmniSchedulerOutput -from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk -from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory -from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) logger = init_logger(__name__) @@ -65,24 +64,9 @@ class OmniARScheduler(VLLMScheduler): # Track requests that have already triggered prefill transfer to avoid duplicates self.transfer_triggered_requests: set[str] = set() model_config = self.vllm_config.model_config - self.omni_connector = None - if model_config.async_chunk: - connector_config = model_config.stage_connector_config - connector_specs = ConnectorSpec( - name=connector_config.get("name", "SharedMemoryConnector"), - extra=connector_config.get("extra", {}), - ) - self.omni_connector = OmniConnectorFactory.create_connector(connector_specs) - - custom_process_next_stage_input_func = getattr( - self.vllm_config.model_config, "custom_process_next_stage_input_func", None - ) - if custom_process_next_stage_input_func: - module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) - module = importlib.import_module(module_path) - self.custom_process_next_stage_input_func = getattr(module, func_name) - - self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None) + self.chunk_transfer_adapter = None + if getattr(model_config, "async_chunk", False): + self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) def _get_kv_transfer_criteria(self) -> dict | None: # Note: vllm_config is available in Scheduler after super().__init__ @@ -152,7 +136,15 @@ class OmniARScheduler(VLLMScheduler): return False def schedule(self) -> SchedulerOutput: # type: ignore[override] - scheduler_output = super().schedule() + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + + try: + scheduler_output = super().schedule() + finally: + if self.chunk_transfer_adapter: + # Add request waiting for chunk to the waiting and running queue + self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) try: # Late import to avoid circulars in some launch modes from .output import OmniNewRequestData @@ -181,9 +173,8 @@ class OmniARScheduler(VLLMScheduler): new_list.append(omni_nr) scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] - if self.omni_connector is not None: - get_chunk(self.omni_connector, scheduler_output) - + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output, self.requests) # Add information about requests needing KV cache transfer finished_reqs = self.get_finished_requests_needing_kv_transfer() except Exception: @@ -312,6 +303,11 @@ class OmniARScheduler(VLLMScheduler): kv_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) + elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK: + # In async chunk mode, request may be in either queue. + # Remove from both to avoid stale queue entries. + stopped_running_reqs.add(request) + stopped_preempted_reqs.add(request) else: stopped_preempted_reqs.add(request) @@ -356,9 +352,8 @@ class OmniARScheduler(VLLMScheduler): num_nans_in_logits=request.num_nans_in_logits, ) ) - if self.omni_connector is not None: - custom_process_next_stage_input_func = self.custom_process_next_stage_input_func - put_chunk(self.omni_connector, pooler_output, request, custom_process_next_stage_input_func) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.save_async(pooler_output, request) else: # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 4b0b93726..684aab9ce 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -16,9 +16,9 @@ from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData -from vllm_omni.distributed.omni_connectors.adapter import get_chunk_for_generation -from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory -from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) from vllm_omni.outputs import OmniModelRunnerOutput @@ -26,15 +26,9 @@ class OmniGenerationScheduler(VLLMScheduler): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) model_config = self.vllm_config.model_config - self.omni_connector = None - if model_config.async_chunk: - connector_config = model_config.stage_connector_config - connector_specs = ConnectorSpec( - name=connector_config.get("name", "SharedMemoryConnector"), - extra=connector_config.get("extra", {}), - ) - self.omni_connector = OmniConnectorFactory.create_connector(connector_specs) - self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None) + self.chunk_transfer_adapter = None + if getattr(model_config, "async_chunk", False): + self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) def schedule(self) -> SchedulerOutput: """Diffusion fast path: @@ -59,18 +53,18 @@ class OmniGenerationScheduler(VLLMScheduler): # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled already_finished_reqs: set[Request] = set() while req_index < len(self.running) and token_budget > 0: request = self.running[req_index] - if self.omni_connector is not None: - get_chunk_for_generation(self.omni_connector, request) - # OMNI: Skip requests that are not in self.requests # This can happen when connector marks request as finished and it's removed from requests if request.request_id not in self.requests or ( - self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED + self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED ): already_finished_reqs.add(request) req_index += 1 @@ -107,12 +101,9 @@ class OmniGenerationScheduler(VLLMScheduler): # independent of pooling_params) while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs: request = self.waiting.peek_request() - if self.omni_connector is not None: - get_chunk_for_generation(self.omni_connector, request) - # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( - self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED + self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED ): # Pop the finished request from waiting queue and don't schedule it self.waiting.pop_request() @@ -154,7 +145,11 @@ class OmniGenerationScheduler(VLLMScheduler): # If fast path scheduled none, fall back to the original scheduling if not num_scheduled_tokens: - return super().schedule() + res = super().schedule() + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) + self.chunk_transfer_adapter.postprocess_scheduler_output(res) + return res # Compute common prefix blocks (aligned with v1) num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) @@ -254,6 +249,11 @@ class OmniGenerationScheduler(VLLMScheduler): new_list.append(omni_nr) scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] + + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) + self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output) + except Exception: # If anything goes wrong, leave the original output unchanged init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") @@ -361,7 +361,7 @@ class OmniGenerationScheduler(VLLMScheduler): # Diffusion request: completes in one step; mark finished and free resources if request.status == RequestStatus.FINISHED_STOPPED or ( - self.omni_connector is None and request.num_computed_tokens >= request.num_prompt_tokens + self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens ): request.status = RequestStatus.FINISHED_STOPPED # Optional: set a stop_reason for front-end clarity @@ -377,6 +377,11 @@ class OmniGenerationScheduler(VLLMScheduler): kv_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) + elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK: + # In async chunk mode, request may be in either queue. + # Remove from both to avoid stale queue entries. + stopped_running_reqs.add(request) + stopped_preempted_reqs.add(request) else: stopped_preempted_reqs.add(request) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e55bbc6f0..cbfa1b1bc 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -298,11 +298,8 @@ class OmniDiffusionConfig: # - Text encoders run on GPU while DiT is on CPU # - DiT runs on GPU while encoders are on CPU enable_cpu_offload: bool = False - # Layer-wise offloading (block-level offloading) parameters enable_layerwise_offload: bool = False - # Number of transformer blocks ready for computation to keep on GPU - layerwise_num_gpu_layers: int = 1 use_fsdp_inference: bool = False pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -335,6 +332,15 @@ class OmniDiffusionConfig: # TODO: do not hard code master_port: int | None = None + # Worker extension class for custom functionality + worker_extension_cls: str | None = None + + # Custom pipeline arguments for custom pipelines + custom_pipeline_args: dict[str, Any] | None = None + + # Diffusion model loading format + diffusion_load_format: str = "default" # "default", "custom_pipeline", "dummy" + # http server endpoint config, would be ignored in local mode host: str | None = None port: int | None = None diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index a421a01ed..f26d89d20 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -77,6 +77,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): mp.set_start_method("spawn", force=True) processes = [] + # Extract worker_extension_cls and custom_pipeline_args from od_config + worker_extension_cls = od_config.worker_extension_cls + custom_pipeline_args = getattr(od_config, "custom_pipeline_args", None) + # Launch all worker processes scheduler_pipe_readers = [] scheduler_pipe_writers = [] @@ -91,6 +95,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): od_config, writer, broadcast_handle, + worker_extension_cls, + custom_pipeline_args, ), name=f"DiffusionWorker-{i}", daemon=True, diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index b61f70b69..7806789b1 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import ( maybe_download_from_modelscope, safetensors_weights_iterator, ) +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -208,13 +209,22 @@ class DiffusersPipelineLoader: allow_patterns_overrides=None, ) - def load_model(self, od_config: OmniDiffusionConfig, load_device: str) -> nn.Module: + def load_model( + self, + od_config: OmniDiffusionConfig, + load_device: str, + load_format: str = "default", + custom_pipeline_name: str | None = None, + ) -> nn.Module: """Load a model with the given configurations.""" target_device = torch.device(load_device) with set_default_torch_dtype(od_config.dtype): with target_device: - model = initialize_model(od_config) - + if load_format == "default": + model = initialize_model(od_config) + elif load_format == "custom_pipeline": + model_cls = resolve_obj_by_qualname(custom_pipeline_name) + model = model_cls(od_config=od_config) logger.debug("Loading weights on %s ...", load_device) # Quantization does not happen in `load_weights` but after it self.load_weights(model) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 839fc1d72..7a476ee73 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -10,7 +10,6 @@ import math from collections.abc import Iterable from dataclasses import dataclass -from typing import Any import numpy as np import torch @@ -1251,6 +1250,39 @@ class Bagel(nn.Module): def prepare_vae_latent(self, curr_kvlens, curr_rope, image_sizes, new_token_ids): return self.prepare_input(curr_kvlens, curr_rope, image_sizes, new_token_ids) + def prepare_vae_latent_cfg(self, curr_kvlens, curr_rope, image_sizes): + packed_position_ids, packed_indexes, packed_key_value_indexes = list(), list(), list() + + query_curr = curr = 0 + for (H, W), curr_kvlen, curr_position_id in zip(image_sizes, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + + h, w = H // self.latent_downsample, W // self.latent_downsample + num_image_tokens = h * w + packed_indexes.extend(range(curr, curr + num_image_tokens)) + curr += num_image_tokens + query_curr += num_image_tokens + + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + + packed_position_ids.extend([curr_position_id] * (num_image_tokens + 2)) + + generation_input = { + "cfg_packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long), + "cfg_key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + "cfg_packed_query_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "cfg_packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + } + + return generation_input + def generate_image( self, packed_text_ids: torch.LongTensor, @@ -1266,11 +1298,25 @@ class Bagel(nn.Module): packed_key_value_indexes: torch.LongTensor, num_timesteps: int = 24, timestep_shift: float = 1.0, + cfg_renorm_min: float = 0.0, + cfg_renorm_type: str = "global", + cfg_interval: tuple[float, float] = [0, 1], + # cfg_text + cfg_text_scale: float = 1.0, + cfg_text_packed_query_indexes: torch.LongTensor | None = None, + cfg_text_packed_position_ids: torch.LongTensor | None = None, + cfg_text_past_key_values: NaiveCache | None = None, + cfg_text_key_values_lens: torch.IntTensor | None = None, + cfg_text_packed_key_value_indexes: torch.LongTensor | None = None, + # cfg_img + cfg_img_scale: float = 1.0, + cfg_img_packed_query_indexes: torch.LongTensor | None = None, + cfg_img_packed_position_ids: torch.LongTensor | None = None, + cfg_img_past_key_values: NaiveCache | None = None, + cfg_img_key_values_lens: torch.IntTensor | None = None, + cfg_img_packed_key_value_indexes: torch.LongTensor | None = None, + cfg_type: str = "parallel", ): - model_pred_cache_dic, model_pred_current = None, None - model_pred_text_cache_dic, model_pred_text_current = None, None - model_pred_img_cache_dic, model_pred_img_current = None, None - x_t = packed_init_noises timesteps = torch.linspace(1, 0, num_timesteps, device=x_t.device) @@ -1280,6 +1326,12 @@ class Bagel(nn.Module): for i, t in enumerate(timesteps): timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device) + if t > cfg_interval[0] and t <= cfg_interval[1]: + cfg_text_scale_ = cfg_text_scale + cfg_img_scale_ = cfg_img_scale + else: + cfg_text_scale_ = 1.0 + cfg_img_scale_ = 1.0 v_t = self._forward_flow( x_t=x_t, timestep=timestep, @@ -1293,17 +1345,26 @@ class Bagel(nn.Module): key_values_lens=key_values_lens, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, - # cache - model_pred_cache_dic=model_pred_cache_dic, - model_pred_current=model_pred_current, - model_pred_text_cache_dic=model_pred_text_cache_dic, - model_pred_text_current=model_pred_text_current, - model_pred_img_cache_dic=model_pred_img_cache_dic, - model_pred_img_current=model_pred_img_current, + cfg_renorm_min=cfg_renorm_min, + cfg_renorm_type=cfg_renorm_type, + # cfg_text + cfg_text_scale=cfg_text_scale_, + cfg_text_packed_position_ids=cfg_text_packed_position_ids, + cfg_text_packed_query_indexes=cfg_text_packed_query_indexes, + cfg_text_key_values_lens=cfg_text_key_values_lens, + cfg_text_past_key_values=cfg_text_past_key_values, + cfg_text_packed_key_value_indexes=cfg_text_packed_key_value_indexes, + # cfg_img + cfg_img_scale=cfg_img_scale_, + cfg_img_packed_position_ids=cfg_img_packed_position_ids, + cfg_img_packed_query_indexes=cfg_img_packed_query_indexes, + cfg_img_key_values_lens=cfg_img_key_values_lens, + cfg_img_past_key_values=cfg_img_past_key_values, + cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes, + cfg_type=cfg_type, ) x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise - unpacked_latent = x_t.split((packed_seqlens - 2).tolist()) return unpacked_latent @@ -1321,13 +1382,23 @@ class Bagel(nn.Module): key_values_lens: torch.IntTensor, past_key_values: NaiveCache, packed_key_value_indexes: torch.LongTensor, - # cache - model_pred_cache_dic: dict[str, Any] | None = None, - model_pred_current: int | None = None, - model_pred_text_cache_dic: dict[str, Any] | None = None, - model_pred_text_current: int | None = None, - model_pred_img_cache_dic: dict[str, Any] | None = None, - model_pred_img_current: int | None = None, + cfg_renorm_min: float = 0.0, + cfg_renorm_type: str = "global", + # cfg_text + cfg_text_scale: float = 1.0, + cfg_text_packed_position_ids: torch.LongTensor | None = None, + cfg_text_packed_query_indexes: torch.LongTensor | None = None, + cfg_text_key_values_lens: torch.Tensor | None = None, + cfg_text_past_key_values: NaiveCache | None = None, + cfg_text_packed_key_value_indexes: torch.LongTensor | None = None, + # cfg_img + cfg_img_scale: float = 1.0, + cfg_img_packed_position_ids: torch.LongTensor | None = None, + cfg_img_packed_query_indexes: torch.LongTensor | None = None, + cfg_img_key_values_lens: torch.Tensor | None = None, + cfg_img_past_key_values: NaiveCache | None = None, + cfg_img_packed_key_value_indexes: torch.LongTensor | None = None, + cfg_type: str = "parallel", ): packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) @@ -1364,4 +1435,70 @@ class Bagel(nn.Module): v_t = self.llm2vae(output.packed_query_sequence) v_t = v_t[packed_vae_token_indexes] + if cfg_text_scale > 1.0: + cfg_text_output = self.language_model.forward( + packed_query_sequence=packed_sequence, + query_lens=packed_seqlens, + packed_query_position_ids=cfg_text_packed_position_ids, + packed_query_indexes=cfg_text_packed_query_indexes, + past_key_values=cfg_text_past_key_values, + key_values_lens=cfg_text_key_values_lens, + packed_key_value_indexes=cfg_text_packed_key_value_indexes, + update_past_key_values=False, + is_causal=False, + **extra_inputs, + ) + cfg_text_v_t = self.llm2vae(cfg_text_output.packed_query_sequence) + cfg_text_v_t = cfg_text_v_t[packed_vae_token_indexes] + + if cfg_img_scale > 1.0: + cfg_img_output = self.language_model.forward( + packed_query_sequence=packed_sequence, + query_lens=packed_seqlens, + packed_query_position_ids=cfg_img_packed_position_ids, + packed_query_indexes=cfg_img_packed_query_indexes, + past_key_values=cfg_img_past_key_values, + key_values_lens=cfg_img_key_values_lens, + packed_key_value_indexes=cfg_img_packed_key_value_indexes, + update_past_key_values=False, + is_causal=False, + **extra_inputs, + ) + cfg_img_v_t = self.llm2vae(cfg_img_output.packed_query_sequence) + cfg_img_v_t = cfg_img_v_t[packed_vae_token_indexes] + + if cfg_text_scale > 1.0: + if cfg_renorm_type == "text_channel": + v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t) + norm_v_t = torch.norm(v_t, dim=-1, keepdim=True) + norm_v_t_text_ = torch.norm(v_t_text_, dim=-1, keepdim=True) + scale = (norm_v_t / (norm_v_t_text_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) + v_t_text = v_t_text_ * scale + if cfg_img_scale > 1.0: + v_t = cfg_img_v_t + cfg_img_scale * (v_t_text - cfg_img_v_t) + else: + v_t = v_t_text + else: + v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t) + + if cfg_img_scale > 1.0: + v_t_ = cfg_img_v_t + cfg_img_scale * (v_t_text_ - cfg_img_v_t) + else: + v_t_ = v_t_text_ + + # NOTE norm is computed over all dimensions, thus currently only supports batch_size = 1 with navit + if cfg_renorm_type == "global": + norm_v_t = torch.norm(v_t) + norm_v_t_ = torch.norm(v_t_) + elif cfg_renorm_type == "channel": + norm_v_t = torch.norm(v_t, dim=-1, keepdim=True) + norm_v_t_ = torch.norm(v_t_, dim=-1, keepdim=True) + else: + raise NotImplementedError(f"{cfg_renorm_type} is not supported") + scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) + v_t = v_t_ * scale + else: + # No CFG + pass + return v_t diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 040f7ddec..340188ba7 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -9,6 +9,7 @@ from __future__ import annotations import json import os from collections.abc import Iterable +from copy import deepcopy from dataclasses import dataclass from math import isqrt @@ -37,7 +38,12 @@ logger = init_logger(__name__) @dataclass class BagelGenParams: num_timesteps: int = 50 - timestep_shift: float = 1.0 + timestep_shift: float = 3.0 + cfg_text_scale: float = 4.0 + cfg_img_scale: float = 1.5 + cfg_interval: tuple = (0.4, 1.0) + cfg_renorm_min: float = 0.0 + cfg_renorm_type: str = "global" def add_special_tokens(tokenizer): @@ -296,10 +302,15 @@ class BagelPipeline(nn.Module): ) image_shape = (height, width) - # Map request params to Bagel gen params (defaults follow Bagel inferencer) + extra_args = getattr(req.sampling_params, "extra_args", {}) or {} + cfg_text_scale = extra_args.get("cfg_text_scale", 4.0) + cfg_img_scale = extra_args.get("cfg_img_scale", 1.5) + gen_params = BagelGenParams( num_timesteps=int(req.sampling_params.num_inference_steps or 50), timestep_shift=3.0, + cfg_text_scale=cfg_text_scale, + cfg_img_scale=cfg_img_scale, ) gen_context = { @@ -307,20 +318,26 @@ class BagelPipeline(nn.Module): "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } + cfg_text_context = deepcopy(gen_context) + cfg_img_context = deepcopy(gen_context) - # Add text prompt (prefill) on gen context. - # [Omni] Check for injected KV Cache from remote transfer injected_kv = req.sampling_params.past_key_values if injected_kv is not None: logger.info("Using injected KV Cache (direct)") gen_context["past_key_values"] = injected_kv - - # User requested: kv_lens and ropes set to [gen_context["past_key_values"].key_cache[0].shape[0]] - # Assuming injected_kv is compatible and has key_cache[0] seq_len = injected_kv.key_cache[0].shape[0] gen_context["kv_lens"] = [seq_len] gen_context["ropes"] = [seq_len] + # Disable CFG: single KV cache cannot support 3-branch CFG + logger.warning("CFG is disabled when using injected KV Cache") + gen_params = BagelGenParams( + num_timesteps=gen_params.num_timesteps, + timestep_shift=gen_params.timestep_shift, + cfg_text_scale=1.0, + cfg_img_scale=1.0, + ) + else: image_input = ( None if isinstance(first_prompt, str) else (first_prompt.get("multi_modal_data") or {}).get("image") @@ -335,9 +352,35 @@ class BagelPipeline(nn.Module): if self.image_processor and self.vae: def vit_transforms(img): - # SigLIP processor returns dict with pixel_values; we want the tensor return self.image_processor(images=img, return_tensors="pt").pixel_values[0] + stride = self.bagel.latent_downsample + max_img_size = int(self.bagel.max_latent_size * stride) + + def _resize_to_stride(img): + if img.mode != "RGB": + img = img.convert("RGB") + w, h = img.size + # Scale down if longest edge exceeds max + scale = min(max_img_size / max(w, h), 1.0) + # Scale up if shortest edge is too small (min 256) + min_img_size = min(256, max_img_size) + scale = max(scale, min_img_size / min(w, h)) + new_w = max(stride, int(round(w * scale / stride) * stride)) + new_h = max(stride, int(round(h * scale / stride) * stride)) + # Clamp to max + new_w = min(new_w, max_img_size) + new_h = min(new_h, max_img_size) + if new_w != w or new_h != h: + img = img.resize((new_w, new_h), Image.BICUBIC) + return img + + image_input = [_resize_to_stride(img) for img in image_input] + + resized_w, resized_h = image_input[0].size + image_shape = (resized_h, resized_w) + logger.info(f"img2img: resized image to {resized_w}x{resized_h}") + def vae_transforms(img): if img.mode != "RGB": img = img.convert("RGB") @@ -345,7 +388,7 @@ class BagelPipeline(nn.Module): arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0 return arr.permute(2, 0, 1) - # 1. Update VAE + # Update gen_context with image (VAE + ViT) gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -353,12 +396,9 @@ class BagelPipeline(nn.Module): transforms=vae_transforms, new_token_ids=self.new_token_ids, ) - for k, v in gen_input_vae.items(): if torch.is_tensor(v): gen_input_vae[k] = v.to(self.device) - - # VAE needs bfloat16 to match model strings usually, specifically encode with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -370,7 +410,6 @@ class BagelPipeline(nn.Module): gen_context["kv_lens"] = newlens_vae gen_context["ropes"] = new_rope_vae - # 2. Update ViT gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -378,11 +417,9 @@ class BagelPipeline(nn.Module): transforms=vit_transforms, new_token_ids=self.new_token_ids, ) - for k, v in gen_input_img.items(): if torch.is_tensor(v): gen_input_img[k] = v.to(self.device) - with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -393,6 +430,10 @@ class BagelPipeline(nn.Module): ) gen_context["kv_lens"] = newlens_img gen_context["ropes"] = new_rope_img + + cfg_text_context = deepcopy(gen_context) + + # Update gen_context with text prompt generation_input, newlens, new_rope = self.bagel.prepare_prompts( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -424,12 +465,56 @@ class BagelPipeline(nn.Module): gen_context["kv_lens"] = newlens gen_context["ropes"] = new_rope + # cfg_text_context: update with negative prompt (no text condition) + neg_prompt = extra_args.get("negative_prompt", "") + neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + prompts=[neg_prompt], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in neg_input.items(): + if torch.is_tensor(v): + neg_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_text_context["past_key_values"], **neg_input + ) + cfg_text_context["kv_lens"] = neg_newlens + cfg_text_context["ropes"] = neg_rope + + # cfg_img_context: update with text prompt (no image condition) + cfg_img_generation_input, cfg_img_newlens, cfg_img_new_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_img_context["kv_lens"], + curr_rope=cfg_img_context["ropes"], + prompts=[prompt], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in cfg_img_generation_input.items(): + if torch.is_tensor(v): + cfg_img_generation_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_img_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_img_context["past_key_values"], **cfg_img_generation_input + ) + cfg_img_context["kv_lens"] = cfg_img_newlens + cfg_img_context["ropes"] = cfg_img_new_rope + if req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.seed) - # Prepare latent query and run flow generation_input = self.bagel.prepare_vae_latent( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], @@ -462,6 +547,25 @@ class BagelPipeline(nn.Module): if torch.is_tensor(v): generation_input[k] = v.to(self.device) + # text cfg + generation_input_cfg_text = self.bagel.prepare_vae_latent_cfg( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + image_sizes=[image_shape], + ) + # img cfg + generation_input_cfg_img = self.bagel.prepare_vae_latent_cfg( + curr_kvlens=cfg_img_context["kv_lens"], + curr_rope=cfg_img_context["ropes"], + image_sizes=[image_shape], + ) + for k, v in generation_input_cfg_text.items(): + if torch.is_tensor(v): + generation_input_cfg_text[k] = v.to(self.device) + for k, v in generation_input_cfg_img.items(): + if torch.is_tensor(v): + generation_input_cfg_img[k] = v.to(self.device) + with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -469,12 +573,26 @@ class BagelPipeline(nn.Module): ): latents = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], + cfg_text_past_key_values=cfg_text_context["past_key_values"], + cfg_img_past_key_values=cfg_img_context["past_key_values"], num_timesteps=gen_params.num_timesteps, timestep_shift=gen_params.timestep_shift, + cfg_text_scale=gen_params.cfg_text_scale, + cfg_img_scale=gen_params.cfg_img_scale, + cfg_interval=gen_params.cfg_interval, + cfg_renorm_min=gen_params.cfg_renorm_min, + cfg_renorm_type=gen_params.cfg_renorm_type, **generation_input, + cfg_text_packed_position_ids=generation_input_cfg_text["cfg_packed_position_ids"], + cfg_text_packed_query_indexes=generation_input_cfg_text["cfg_packed_query_indexes"], + cfg_text_key_values_lens=generation_input_cfg_text["cfg_key_values_lens"], + cfg_text_packed_key_value_indexes=generation_input_cfg_text["cfg_packed_key_value_indexes"], + cfg_img_packed_position_ids=generation_input_cfg_img["cfg_packed_position_ids"], + cfg_img_packed_query_indexes=generation_input_cfg_img["cfg_packed_query_indexes"], + cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"], + cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"], ) - # Decode first sample img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) return DiffusionOutput(output=img) diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py index e60bcbe5a..071716657 100644 --- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py +++ b/vllm_omni/diffusion/models/sd3/sd3_transformer.py @@ -2,14 +2,20 @@ from collections.abc import Iterable import torch import torch.nn as nn -from diffusers.models.attention import FeedForward +import torch.nn.functional as F # TODO replace this with vLLM implementation from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed from diffusers.models.modeling_outputs import Transformer2DModelOutput -from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm, SD35AdaLayerNormZeroX +from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, SD35AdaLayerNormZeroX from vllm.logger import init_logger -from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention @@ -18,6 +24,54 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig logger = init_logger(__name__) +class GELU(nn.Module): + def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): + super().__init__() + self.proj = ColumnParallelLinear(dim_in, dim_out, bias=bias) + self.approximate = approximate + + def forward(self, hidden_states): + hidden_states, _ = self.proj(hidden_states) + hidden_states = F.gelu(hidden_states, approximate=self.approximate) + return hidden_states + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: int | None = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "gelu-approximate", + final_dropout: bool = False, + inner_dim=None, + bias: bool = True, + ): + super().__init__() + if inner_dim is None: + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == "gelu-approximate": + act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias) + else: + raise ValueError(f"Unsupported activation function type: {activation_fn}") + + self.net = nn.ModuleList([]) + self.net.append(act_fn) + self.net.append(nn.Dropout(dropout)) + self.net.append(RowParallelLinear(inner_dim, dim_out, bias=bias)) + if final_dropout: + self.net.append(nn.Dropout(dropout)) + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + for layer in self.net: + output = layer(hidden_states) + hidden_states = output[0] if isinstance(output, tuple) else output + return hidden_states + + class SD3PatchEmbed(nn.Module): """ 2D Image to Patch Embedding with support for SD3. @@ -77,7 +131,6 @@ class SD3CrossAttention(nn.Module): hidden_size=dim, head_size=self.head_dim, total_num_heads=num_heads, - disable_tp=True, ) self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() @@ -88,17 +141,18 @@ class SD3CrossAttention(nn.Module): added_kv_proj_dim, head_size=self.inner_kv_dim // self.num_heads, total_num_heads=self.num_heads, - disable_tp=True, ) + else: + self.add_kv_proj = None if not context_pre_only: - self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias) + self.to_add_out = RowParallelLinear(self.inner_dim, self.dim, bias=out_bias) else: self.to_add_out = None if not pre_only: self.to_out = nn.ModuleList([]) - self.to_out.append(ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)) + self.to_out.append(RowParallelLinear(self.inner_dim, self.dim, bias=out_bias)) else: self.to_out = None @@ -106,7 +160,7 @@ class SD3CrossAttention(nn.Module): self.norm_added_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.attn = Attention( - num_heads=num_heads, + num_heads=self.to_qkv.num_heads, head_size=self.head_dim, softmax_scale=1.0 / (self.head_dim**0.5), causal=False, @@ -118,13 +172,15 @@ class SD3CrossAttention(nn.Module): encoder_hidden_states: torch.Tensor | None = None, ): # Compute QKV for image stream (sample projections) - qkv, _ = self.to_qkv(hidden_states) + qkv = self.to_qkv(hidden_states) + qkv = qkv[0] img_query, img_key, img_value = qkv.chunk(3, dim=-1) # Reshape for multi-head attention - img_query = img_query.unflatten(-1, (self.num_heads, -1)) - img_key = img_key.unflatten(-1, (self.num_heads, -1)) - img_value = img_value.unflatten(-1, (self.num_heads, -1)) + local_num_heads = self.to_qkv.num_heads + img_query = img_query.unflatten(-1, (local_num_heads, -1)) + img_key = img_key.unflatten(-1, (local_num_heads, -1)) + img_value = img_value.unflatten(-1, (local_num_heads, -1)) # Apply QK normalization img_query = self.norm_q(img_query) @@ -132,12 +188,13 @@ class SD3CrossAttention(nn.Module): if encoder_hidden_states is not None: # Compute QKV for text stream (context projections) - qkv, _ = self.add_kv_proj(encoder_hidden_states) - txt_query, txt_key, txt_value = qkv.chunk(3, dim=-1) + qkv_add = self.add_kv_proj(encoder_hidden_states) + qkv_add = qkv_add[0] + txt_query, txt_key, txt_value = qkv_add.chunk(3, dim=-1) - txt_query = txt_query.unflatten(-1, (self.num_heads, -1)) - txt_key = txt_key.unflatten(-1, (self.num_heads, -1)) - txt_value = txt_value.unflatten(-1, (self.num_heads, -1)) + txt_query = txt_query.unflatten(-1, (local_num_heads, -1)) + txt_key = txt_key.unflatten(-1, (local_num_heads, -1)) + txt_value = txt_value.unflatten(-1, (local_num_heads, -1)) txt_query = self.norm_added_q(txt_query) txt_key = self.norm_added_k(txt_key) @@ -265,6 +322,8 @@ class SD3TransformerBlock(nn.Module): encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor, ) -> tuple[torch.Tensor, torch.Tensor]: + if isinstance(encoder_hidden_states, tuple): + encoder_hidden_states = encoder_hidden_states[0] if self.use_dual_attention: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1( hidden_states, emb=temb @@ -363,7 +422,7 @@ class SD3Transformer2DModel(nn.Module): self.time_text_embed = CombinedTimestepTextProjEmbeddings( embedding_dim=self.inner_dim, pooled_projection_dim=self.pooled_projection_dim ) - self.context_embedder = nn.Linear(self.joint_attention_dim, self.caption_projection_dim) + self.context_embedder = ReplicatedLinear(self.joint_attention_dim, self.caption_projection_dim) self.transformer_blocks = nn.ModuleList( [ @@ -380,7 +439,9 @@ class SD3Transformer2DModel(nn.Module): ) self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) - self.proj_out = nn.Linear(self.inner_dim, self.patch_size * self.patch_size * self.out_channels, bias=True) + self.proj_out = ReplicatedLinear( + self.inner_dim, self.patch_size * self.patch_size * self.out_channels, bias=True + ) def forward( self, @@ -427,6 +488,9 @@ class SD3Transformer2DModel(nn.Module): hidden_states = self.norm_out(hidden_states, temb) hidden_states = self.proj_out(hidden_states) + if isinstance(hidden_states, (tuple, list)): + hidden_states = hidden_states[0] + # unpatchify patch_size = self.patch_size height = height // patch_size diff --git a/vllm_omni/diffusion/offload.py b/vllm_omni/diffusion/offload.py deleted file mode 100644 index 6f1d8a0db..000000000 --- a/vllm_omni/diffusion/offload.py +++ /dev/null @@ -1,535 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU offloading utilities for diffusion models. - -This module provides mutual-exclusion CPU offloading between DiT and encoders. -When enable_cpu_offload is enabled: -- Text encoders run on GPU while DiT is on CPU -- DiT runs on GPU while encoders are offloaded to CPU - -This allows running large models on limited GPU memory. -""" - -from __future__ import annotations - -from functools import partial -from itertools import chain -from typing import TYPE_CHECKING, Any - -import torch -from torch import nn -from vllm.logger import init_logger - -from vllm_omni.platforms import current_omni_platform - -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - -logger = init_logger(__name__) - - -class SequentialOffloader: - """Sequential offloader: DiT and encoders take turns on GPU. - - Uses PyTorch's forward pre-hooks to automatically swap models: - - Before encoder runs: move DiT modules to CPU, move encoder to GPU - - Before DiT runs: move encoders to CPU, move active DiT to GPU - - This ensures only one large model group is on GPU at a time. - """ - - def __init__( - self, - dits: list[nn.Module], - encoders: list[nn.Module], - device: torch.device, - pin_memory: bool = True, - ): - assert all(isinstance(m, nn.Module) for m in dits), "All dits must be nn.Module" - assert all(isinstance(m, nn.Module) for m in encoders), "All encoders must be nn.Module" - self.dits = dits - self.encoders = encoders - self.device = device - self.pin_memory = pin_memory - self._handles: list = [] - - def _to_cpu(self, module: nn.Module) -> None: - """Move module to CPU with optional memory pinning.""" - # Skip if already on CPU - try: - param = next(module.parameters()) - if param.device.type == "cpu": - return - except StopIteration: - return - - previous_device = param.device - module.to("cpu", non_blocking=True) - - # Release allocator blocks when tensors leave the GPU. - if previous_device.type != "cpu": - torch.cuda.empty_cache() - - if self.pin_memory: - for p in module.parameters(): - if p.data.device.type == "cpu" and not p.data.is_pinned(): - p.data = p.data.pin_memory() - - def _to_gpu(self, module: nn.Module) -> None: - """Move module to GPU.""" - # Skip if already on target device - try: - if next(module.parameters()).device == self.device: - return - except StopIteration: - return - - module.to(self.device, non_blocking=True) - - def _dit_pre_hook(self, module: nn.Module, args: tuple) -> None: - """Before DiT forward: offload encoders, load DiT.""" - for enc in self.encoders: - self._to_cpu(enc) - self._to_gpu(module) - - current_omni_platform.synchronize() - - logger.debug("Swapped: encoders -> CPU, DiT -> GPU") - - def _encoder_pre_hook(self, module: nn.Module, args: tuple) -> None: - """Before encoder forward: offload DiT, load encoder.""" - for dit_mod in self.dits: - self._to_cpu(dit_mod) - self._to_gpu(module) - - current_omni_platform.synchronize() - - logger.debug("Swapped: DiT -> CPU, encoder -> GPU") - - def register(self) -> None: - """Register forward pre-hooks on DiT and encoders.""" - # Hook on each DiT-like module - for dit_mod in self.dits: - h = dit_mod.register_forward_pre_hook(self._dit_pre_hook) - self._handles.append(h) - logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) - - # Hook on each encoder - for enc in self.encoders: - h = enc.register_forward_pre_hook(self._encoder_pre_hook) - self._handles.append(h) - logger.debug("Registered offload hook for %s", enc.__class__.__name__) - - def remove(self) -> None: - """Remove all hooks.""" - for h in self._handles: - h.remove() - self._handles = [] - - -class LayerwiseOffloader: - """Layer-wise CPU offloading for transformer blocks. - - Keeps only a sliding window of layers (blocks), by default a single layer, on GPU, - prefetching the next block while the current block computes to approach compute - memcpy overlap. - Unused blocks are freed on GPU. - - Based on implementations from: - https://github.com/sgl-project/sglang/blob/v0.5.8/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py - """ - - def __init__( - self, - blocks: list[nn.Module], - device: torch.device, - pin_memory: bool = True, - num_gpu_layers: int = 1, - ): - assert all(isinstance(m, nn.Module) for m in blocks), "All transformer blocks must be torch.nn.Module" - assert current_omni_platform.is_cuda(), "Layerwise offloading is only supported on cuda devices for now" - - self.blocks = blocks - self.device = device - self.pin_memory = pin_memory - self.num_gpu_layers = num_gpu_layers - self.num_blocks = len(self.blocks) - if self.num_blocks == 0: - raise ValueError("LayerwiseOffloader requires at least one block, but found 0.") - if not (1 <= self.num_gpu_layers <= self.num_blocks): - raise ValueError(f"Invalid num_gpu_layers {self.num_gpu_layers} with {self.num_blocks} blocks") - - self._pre_hook_handles: list = [] - self._post_hook_handles: list = [] - - self._copy_stream = torch.cuda.Stream() - - # Per-layer synchronization primitive: set after H2D copy completes. - self._prefetch_done: list[torch.cuda.Event | None] = [None] * self.num_blocks - - # Simple state to avoid redundant work. - self._resident: list[bool] = [False] * self.num_blocks - - # Pre-allocate gpu tensors - # layer-id -> {dtype -> flattened aggregated cpu tensor} - self.layer_cpu_weights: list[dict[torch.dtype, torch.Tensor]] = [] - self.layer_metadata: list[dict[torch.dtype, list[dict[str, Any]]]] = [] - - self.block_parameters: dict[int, dict[str, nn.Parameter]] = {} - self.block_buffers: dict[int, dict[str, torch.Tensor]] = {} - for layer_idx, block in enumerate(self.blocks): - self.block_parameters[layer_idx] = dict(block.named_parameters()) - self.block_buffers[layer_idx] = dict(block.named_buffers()) - - dtype_cpu_flattened_weights, dtype_metadata = self._to_cpu( - self.block_parameters[layer_idx], self.block_buffers[layer_idx] - ) - self.layer_cpu_weights.append(dtype_cpu_flattened_weights) - self.layer_metadata.append(dtype_metadata) - - if self.num_blocks != len(self.layer_cpu_weights): - logger.error( - f"Inconsistent block layers happened: # of blocks: {self.num_blocks}; " - f"# of layer cpu weights: {len(self.layer_cpu_weights)}" - ) - - # Register pre and post forward hooks on each of the blocks - self.register_block_hooks() - - # Pre-fetch the first layer - # For subsequent requests, the first layer/block will be pre-fetched - # during the last layer compute of the previous request. - self.prefetch_layer(0, non_blocking=False) - - def _to_cpu( - self, params: dict[str, nn.Parameter], bufs: dict[str, torch.Tensor] - ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: - """Move block parameters and buffers to CPU, flattening by dtype. - - Consolidates parameters and buffers into contiguous CPU tensors grouped by dtype - for GPU transfers. Replaces original tensors with empty placeholders. - - Returns: - Tuple of - flattened CPU tensors by dtype, - metadata for reconstruction by dtype - """ - dtype_grouped_weights: dict[torch.dtype, dict[str, torch.Tensor]] = {} - dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {} - # order does matter - dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} - - for name, param_or_buf in chain(params.items(), bufs.items()): - dtype = param_or_buf.dtype - if dtype not in dtype_grouped_weights: - dtype_grouped_weights[dtype] = {} - dtype_grouped_weights[dtype][name] = param_or_buf - - for dtype, name2weights in dtype_grouped_weights.items(): - # total # of parameters + buffers - total_numel = sum(t.numel() for _, t in name2weights.items()) - cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=self.pin_memory) - - current_offset = 0 - for name, param_or_buf in name2weights.items(): - numel = param_or_buf.numel() - cpu_tensor[current_offset : current_offset + numel].copy_(param_or_buf.flatten()) - if dtype not in dtype_metadata: - dtype_metadata[dtype] = [] - dtype_metadata[dtype].append( - { - "name": name, - "offset": current_offset, - "numel": numel, - "shape": param_or_buf.shape, - } - ) - - param_or_buf.data = torch.empty((), device=self.device, dtype=dtype) - current_offset += numel - - dtype_cpu_flattened_weights[dtype] = cpu_tensor - - return dtype_cpu_flattened_weights, dtype_metadata - - def register_block_hooks(self) -> None: - """Register forward hooks on blocks for prefetching and offloading.""" - - def _pre_hook(module: nn.Module, args: tuple, *, layer_idx: int) -> None: - # For the last block / layer, prefetch layer 0 (the first layer) - next_id = (layer_idx + 1) % self.num_blocks - self.prefetch_layer(next_id, non_blocking=True) - - def _post_hook(module: nn.Module, args: tuple, output: tuple, *, layer_idx: int) -> None: - self.offload_layer(layer_idx) - self._resident[layer_idx] = False - self._prefetch_done[layer_idx] = None - - for i, layer in enumerate(self.blocks): - pre_hook_fn = partial(_pre_hook, layer_idx=i) - handle = layer.register_forward_pre_hook(pre_hook_fn) - self._pre_hook_handles.append(handle) - - post_hook_fn = partial(_post_hook, layer_idx=i) - handle = layer.register_forward_hook(post_hook_fn) - self._post_hook_handles.append(handle) - - @torch.compiler.disable - def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None: - """Copy layer weights from CPU -> GPU. - - Pre-fetch target layer in an asynchronous way with compute - memory copy overlap, - with non_blocking set to True. - """ - if layer_idx >= self.num_blocks or layer_idx < 0: - logger.warning(f"Invalid layer id specified: {layer_idx}") - return - - self._copy_stream.wait_stream(torch.cuda.current_stream()) - - layers_to_fetch = [(layer_idx + i) % self.num_blocks for i in range(self.num_gpu_layers)] - - for idx in layers_to_fetch: - if self._resident[idx]: - continue - - layer_params = self.block_parameters[idx] - layer_bufs = self.block_buffers[idx] - - evt = torch.cuda.Event() - gpu_weights: dict[torch.dtype, torch.Tensor] = {} - - with torch.cuda.stream(self._copy_stream): - for dtype, cpu_weight in self.layer_cpu_weights[idx].items(): - gpu_weight = torch.empty(cpu_weight.shape, dtype=dtype, device=self.device) - gpu_weight.copy_(cpu_weight, non_blocking=non_blocking) - gpu_weights[dtype] = gpu_weight - - evt.record(self._copy_stream) - - for dtype in self.layer_metadata[idx]: - ordered_metadata: list[dict[str, Any]] = self.layer_metadata[idx][dtype] - - gpu_weight = gpu_weights[dtype] - - for metadata in ordered_metadata: - target_name = metadata["name"] - target_param_or_buf = ( - layer_params[target_name] if target_name in layer_params else layer_bufs[target_name] - ) - - target_param_or_buf.data = gpu_weight[ - metadata["offset"] : metadata["offset"] + metadata["numel"] - ].view(metadata["shape"]) - - self._prefetch_done[idx] = evt - self._resident[idx] = True - - @torch.compiler.disable - def offload_layer(self, layer_idx: int) -> None: - """Free GPU memory for layer by replacing tensors with empty placeholders.""" - if layer_idx >= self.num_blocks or layer_idx < 0: - logger.warning(f"Invalid layer id specified: {layer_idx}") - return - if not self._resident[layer_idx]: - logger.warning(f"{layer_idx} is not residing on GPU") - return - - evt = self._prefetch_done[layer_idx] - if evt is not None: - torch.cuda.current_stream().wait_event(evt) - - # free GPU residency - for _, param in self.block_parameters[layer_idx].items(): - param.data = torch.empty((), device=self.device, dtype=param.dtype) - for _, buf in self.block_buffers[layer_idx].items(): - buf.data = torch.empty((), device=self.device, dtype=buf.dtype) - - def remove_all_hooks(self) -> None: - """Remove all hooks.""" - for h in self._pre_hook_handles: - h.remove() - for h in self._post_hook_handles: - h.remove() - self._pre_hook_handles.clear() - self._post_hook_handles.clear() - - @staticmethod - def get_blocks_attr_name(model: nn.Module) -> str | None: - """Retrieve blocks attribute name from provided DiT model""" - return getattr(model.__class__, "_layerwise_offload_blocks_attr", None) - - @staticmethod - def get_blocks_from_dit(model: nn.Module) -> list[nn.Module]: - """ - Retrieve a list of blocks from provided DiT model. Blocks attribute name - are found by `_layerwise_offload_blocks_attr` set to DiT models. For example, - - ``` - class WanTransformer3DModel(nn.Module): - _layerwise_offload_blocks_attr = "blocks" - ``` - """ - blocks_attr_name = LayerwiseOffloader.get_blocks_attr_name(model) - if blocks_attr_name is None: - logger.warning( - f"No _layerwise_offload_blocks_attr defined for {model.__class__.__name__}, " - "skipping layerwise offloading" - ) - return [] - - _blocks = getattr(model, blocks_attr_name, None) - if _blocks is None: - logger.warning( - f"Blocks (layers) '{blocks_attr_name}' not found on {model.__class__.__name__}, " - "skipping layerwise offloading" - ) - return [] - - return list(_blocks) - - -def apply_offload_hooks( - model: nn.Module, - od_config: OmniDiffusionConfig, - *, - device: torch.device | None = None, -) -> None: - """Apply mutual-exclusion offload hooks based on config. - - When enable_cpu_offload is enabled, DiT and encoders swap GPU access: - - Encoders (text_encoder, text_encoder_2, text_encoder_3, image_encoder) - run on GPU while DiT is on CPU - - DiT runs on GPU while encoders are on CPU - - Args: - model: Diffusion pipeline model - od_config: OmniDiffusionConfig with offload settings - """ - enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False) - enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False) - pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) - - if not enable_cpu_offload and not enable_layerwise_offload: - return - if enable_cpu_offload and enable_layerwise_offload: - # NOTE: Model-wise and layerwise cpu offloading are not supported together at this moment, - # consider layerwise offloading has higher priority than model-wise offloading - enable_cpu_offload = False - logger.info( - "Model-wise and layer-wise CPU offloading are not supported together at this moment. " - "Automatically disabled model-wise offloading." - ) - # For now, model-wise and layer-wise (block-wise) offloading - # are functioning as expected when cuda device is available - if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 1: - logger.info("CPU Offloading requires cuda devices available. Skipping for now...") - return - - # Find DiT/transformer modules - dit_modules: list[nn.Module] = [] - dit_names: list[str] = [] - candidate_attrs = ["transformer", "transformer_2", "dit"] - for attr in candidate_attrs: - if not hasattr(model, attr): - continue - module_obj = getattr(model, attr) - if module_obj is None: - continue - - assert isinstance(module_obj, nn.Module), f"Expected {attr} to be nn.Module, got {type(module_obj)!r}" - - if module_obj in dit_modules: - continue - - dit_modules.append(module_obj) - dit_names.append(attr) - - if not dit_modules: - logger.warning("enable_cpu_offload enabled but no transformer/dit/unet found") - return - if device is None: - try: - device = next(dit_modules[0].parameters()).device - except StopIteration: - try: - device = current_omni_platform.get_torch_device() - except (NotImplementedError, AttributeError): - logger.error("Fail to get device of pipeline. Skipping applying offloading hooks") - return - - # Collect all encoders - encoders: list[nn.Module] = [] - encoder_names: list[str] = [] - for attr in ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]: - if hasattr(model, attr) and getattr(model, attr) is not None: - encoders.append(getattr(model, attr)) - encoder_names.append(attr) - if not encoders and enable_cpu_offload: - logger.warning("enable_cpu_offload enabled but no encoders found") - return - for enc in encoders: - enc.to(device) - - # Collect VAE - for name in ["vae"]: - module = getattr(model, name, None) - if module is None: - continue - try: - module.to(device, non_blocking=True) - except Exception as exc: - logger.debug("Failed to move %s to GPU: %s", name, exc) - - if enable_cpu_offload: - # Initial state: keep DiT modules on CPU (encoders typically run first) - for dit_mod in dit_modules: - dit_mod.to("cpu") - - torch.cuda.empty_cache() - - if pin_cpu_memory: - for dit_mod in dit_modules: - for p in dit_mod.parameters(): - if p.data.device.type == "cpu" and not p.data.is_pinned(): - p.data = p.data.pin_memory() - - # Register sequential offload hooks - SequentialOffloader(dit_modules, encoders, device, pin_cpu_memory).register() - logger.info( - "CPU offload enabled: %s <-> %s (mutual exclusion)", - ", ".join(dit_names), - ", ".join(encoder_names), - ) - elif enable_layerwise_offload: - logger.info(f"Applying offloading hooks on {dit_names}") - - for i, dit_module in enumerate(dit_modules): - logger.info(f"Applying hook on {dit_names[i]} ({dit_module.__class__.__name__})") - blocks_attr_name = LayerwiseOffloader.get_blocks_attr_name(dit_module) - blocks = LayerwiseOffloader.get_blocks_from_dit(dit_module) - - if not blocks_attr_name or not blocks: - logger.warning( - "Target layers (blocks) are not found. " - f"Skipping offloading on {dit_names[i]} ({dit_module.__class__.__name__})" - ) - continue - - # move modules other than blocks to gpu and keep them on gpu - for name, m in dit_module.named_children(): - # Skip the blocks module (layers to be offloaded) - if name == blocks_attr_name: - logger.debug(f"Skipped module {name}") - continue - - m.to(device) - logger.debug(f"Moved {name} to device {device}") - - # set to the module (transformer) - offloader = LayerwiseOffloader(blocks, device, pin_cpu_memory, od_config.layerwise_num_gpu_layers) - setattr(dit_module, "_layerwise_offloader", offloader) - - logger.info( - f"Layerwise offloading enabled on {len(blocks)} layers (blocks), " - f"with {od_config.layerwise_num_gpu_layers} kept on device" - ) diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py new file mode 100644 index 000000000..f4fd78c7d --- /dev/null +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.platforms import current_omni_platform + +from .base import OffloadBackend, OffloadConfig, OffloadStrategy +from .layerwise_backend import LayerWiseOffloadBackend +from .sequential_backend import ( + ModelLevelOffloadBackend, + apply_sequential_offload, + remove_sequential_offload, +) + +logger = init_logger(__name__) + +__all__ = [ + "OffloadBackend", + "OffloadConfig", + "OffloadStrategy", + "LayerWiseOffloadBackend", + "ModelLevelOffloadBackend", + "apply_sequential_offload", + "remove_sequential_offload", + "get_offload_backend", +] + + +def get_offload_backend( + od_config: OmniDiffusionConfig, + device: torch.device | None = None, +) -> OffloadBackend | None: + """Create appropriate offload backend based on configuration. + + Args: + od_config: OmniDiffusionConfig with offload settings + device: Target device (auto-detected if None) + + Returns: + OffloadBackend instance or None if offloading disabled + + Example: + >>> backend = get_offload_backend(od_config, device=torch.device("cuda:0")) + >>> if backend: + ... backend.enable(pipeline) + """ + # Extract and validate configuration + config = OffloadConfig.from_od_config(od_config) + + # Return None if no offloading requested + if config.strategy == OffloadStrategy.NONE: + return None + + # Validate platform (CUDA required for now) + if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 1: + logger.warning("CPU offloading requires CUDA devices. Skipping offloading.") + return None + + # Detect device if not provided + if device is None: + try: + device = current_omni_platform.get_torch_device() + except (NotImplementedError, AttributeError) as exc: + logger.error("Failed to detect device: %s. Skipping offloading.", exc) + return None + + # Create appropriate backend + if config.strategy == OffloadStrategy.MODEL_LEVEL: + return ModelLevelOffloadBackend(config, device) + elif config.strategy == OffloadStrategy.LAYER_WISE: + return LayerWiseOffloadBackend(config, device) + else: + logger.error("Unknown offload strategy: %s", config.strategy) + return None diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py new file mode 100644 index 000000000..6f838d4dd --- /dev/null +++ b/vllm_omni/diffusion/offloader/base.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + + +class OffloadStrategy(Enum): + NONE = "none" + MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders + LAYER_WISE = "layer_wise" # Block-level + + +@dataclass +class OffloadConfig: + strategy: OffloadStrategy + pin_cpu_memory: bool = True + + @classmethod + def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": + """Extract and validate offload settings from OmniDiffusionConfig. + + For now, enforces mutual exclusion between model-level and layer-wise offloading. + Layer-wise takes priority if both are enabled. + + Args: + od_config: OmniDiffusionConfig with offload settings + + Returns: + OffloadConfig with validated settings + """ + enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False) + enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False) + pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) + + # Determine strategy (mutual exclusion, layer-wise takes priority) + if enable_layerwise_offload: + strategy = OffloadStrategy.LAYER_WISE + if enable_cpu_offload: + logger.info( + "Both model-level and layer-wise offloading enabled. " + "Layer-wise takes priority, disabling model-level offloading." + ) + elif enable_cpu_offload: + strategy = OffloadStrategy.MODEL_LEVEL + else: + strategy = OffloadStrategy.NONE + + return cls( + strategy=strategy, + pin_cpu_memory=pin_cpu_memory, + ) + + +class OffloadBackend(ABC): + """Base class for CPU offload backends""" + + def __init__(self, config: OffloadConfig, device: torch.device): + self.config = config + self.device = device + self.enabled = False + + @abstractmethod + def enable(self, pipeline: nn.Module) -> None: + """Enable offloading on the pipeline. + + Discovers modules, moves them to appropriate devices, and + registers forward hooks for swapping/prefetching. + + Args: + pipeline: Diffusion pipeline model (e.g., Wan22Pipeline) + """ + raise NotImplementedError + + @abstractmethod + def disable(self) -> None: + """Disable offloading and cleanup resources. + + Removes all registered hooks. Does NOT move modules back to + original devices (caller responsible for that). + """ + raise NotImplementedError + + def is_enabled(self) -> bool: + return self.enabled diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py new file mode 100644 index 000000000..a69b13e64 --- /dev/null +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import chain +from typing import Any + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.platforms import current_omni_platform + +from .base import OffloadBackend, OffloadConfig +from .module_collector import ModuleDiscovery + +logger = init_logger(__name__) + + +class LayerwiseOffloadHook(ModelHook): + """Hook for layerwise (transformer-block-wise) CPU offloading. + + The hook instance retains parameters for both the current registered block + module and those for the next block, as well as flattened CPU tensors which + record the parameters of the current block module, so that these parameters + could be re-materialized on device in an overlapping way. + This hook should be registered to each of the transformer blocks in DiT + module(s) of the target pipeline. + + Based on implementations from: + https://github.com/sgl-project/sglang/blob/v0.5.8/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py + """ + + _HOOK_NAME = "layerwise_offload" + + def __init__( + self, + next_block: nn.Module, + device: torch.device, + stream: torch.cuda.Stream | None = None, + pin_memory: bool = True, + ): + assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" + assert current_omni_platform.is_cuda(), "Layerwise offloading is only supported on cuda devices for now" + + self.next_block = next_block + self.device = device + self.copy_stream = stream or torch.cuda.current_stream() + self.pin_memory = pin_memory + + # Per-block synchronization primitive: set after H2D copy completes. + self._prefetch_done: torch.cuda.Event | None = None + + self.next_block_parameters: dict[str, nn.Parameter] = {} + self.next_block_buffers: dict[str, torch.Tensor] = {} + self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {} + self.dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + + def initialize_hook(self, module: nn.Module) -> nn.Module: + # This all happen during the hook instance being registered to hook registry; + # the input module is kept intact + module = super().initialize_hook(module) + + self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters()) + self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers()) + + self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters()) + self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers()) + + # Pre-allocate gpu tensors in a flattened way + self.dtype_cpu_flattened_weights, self.dtype_metadata = LayerwiseOffloadHook._to_cpu( + self.next_block_parameters, self.next_block_buffers, self.device, self.pin_memory + ) + + return module + + @staticmethod + def _to_cpu( + params: dict[str, nn.Parameter], + bufs: dict[str, torch.Tensor], + device: torch.device, + pin_memory: bool = True, + ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: + """Helper method to move block parameters and buffers to CPU, flattening by dtype. + + Consolidates parameters and buffers into contiguous CPU tensors grouped by dtype + for GPU transfers. Replaces original tensors with empty placeholders. + + Returns: + Tuple of + flattened CPU tensors by dtype, + metadata for reconstruction by dtype + """ + dtype_grouped_weights: dict[torch.dtype, dict[str, torch.Tensor]] = {} + dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {} + # NOTE: order does matter + dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + + for name, param_or_buf in chain(params.items(), bufs.items()): + dtype = param_or_buf.dtype + if dtype not in dtype_grouped_weights: + dtype_grouped_weights[dtype] = {} + dtype_grouped_weights[dtype][name] = param_or_buf + + for dtype, name2weights in dtype_grouped_weights.items(): + # total # of parameters + buffers + total_numel = sum(t.numel() for _, t in name2weights.items()) + cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory) + + current_offset = 0 + for name, param_or_buf in name2weights.items(): + numel = param_or_buf.numel() + cpu_tensor[current_offset : current_offset + numel].copy_(param_or_buf.flatten()) + if dtype not in dtype_metadata: + dtype_metadata[dtype] = [] + dtype_metadata[dtype].append( + { + "name": name, + "offset": current_offset, + "numel": numel, + "shape": param_or_buf.shape, + } + ) + + param_or_buf.data = torch.empty((), device=device, dtype=dtype) + current_offset += numel + + dtype_cpu_flattened_weights[dtype] = cpu_tensor + + return dtype_cpu_flattened_weights, dtype_metadata + + @torch.compiler.disable + def prefetch_layer(self, non_blocking: bool = True) -> None: + """Copy layer weights from CPU -> GPU. + + Pre-fetch target block in an asynchronous way with compute - memory copy overlap, + with non_blocking set to True. + """ + self.copy_stream.wait_stream(torch.cuda.current_stream()) + + layer_params = self.next_block_parameters + layer_bufs = self.next_block_buffers + + evt = torch.cuda.Event() + gpu_weights: dict[torch.dtype, torch.Tensor] = {} + + with torch.cuda.stream(self.copy_stream): + for dtype, cpu_weight in self.dtype_cpu_flattened_weights.items(): + gpu_weight = torch.empty(cpu_weight.shape, dtype=dtype, device=self.device) + gpu_weight.copy_(cpu_weight, non_blocking=non_blocking) + gpu_weights[dtype] = gpu_weight + + evt.record(self.copy_stream) + + for dtype, ordered_metadata in self.dtype_metadata.items(): + # ordered_metadata: list[dict[str, Any]] + gpu_weight = gpu_weights[dtype] + + for metadata in ordered_metadata: + target_name = metadata["name"] + target_param_or_buf = ( + layer_params[target_name] if target_name in layer_params else layer_bufs[target_name] + ) + + target_param_or_buf.data = gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view( + metadata["shape"] + ) + + self._prefetch_done = evt + + @torch.compiler.disable + def offload_layer(self) -> None: + """Free GPU memory for layer by replacing tensors with empty placeholders. + This function does not actually offload weights from GPU back to CPU. + """ + evt = self._prefetch_done + if evt is not None: + torch.cuda.current_stream().wait_event(evt) + + self._prefetch_done = None + + # free GPU residency + for _, param in self.block_parameters.items(): + param.data = torch.empty((), device=self.device, dtype=param.dtype) + for _, buf in self.block_buffers.items(): + buf.data = torch.empty((), device=self.device, dtype=buf.dtype) + + def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: + self.prefetch_layer(non_blocking=True) + + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + self.offload_layer() + + return output + + +def apply_block_hook( + module: nn.Module, + next_block: nn.Module, + device: torch.device, + stream: torch.cuda.Stream | None = None, + pin_memory: bool = True, +) -> LayerwiseOffloadHook: + registry = HookRegistry.get_or_create(module) + hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory) + registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook) + + return hook + + +def remove_block_hook(module: nn.Module) -> None: + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is not None: + registry.remove_hook(LayerwiseOffloadHook._HOOK_NAME) + logger.debug("Removed offload hook from %s", module.__class__.__name__) + + +class LayerWiseOffloadBackend(OffloadBackend): + """Layer-wise (block-level) offloading backend. + + Implements sliding window offloading where only a small number of transformer + blocks reside on GPU at a time. Blocks are prefetched asynchronously while + previous blocks compute, and freed after use. + """ + + def __init__(self, config: OffloadConfig, device: torch.device): + super().__init__(config, device) + + self.copy_stream = torch.cuda.Stream() + self._blocks: list[list[nn.Module]] = [] + + def enable(self, pipeline: nn.Module) -> None: + if self.enabled: + logger.warning("LayerWiseOffloadBackend already enabled") + return + + modules = ModuleDiscovery.discover(pipeline) + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping layer-wise offloading") + return + + # Move encoders to GPU (they stay resident) + for enc in modules.encoders: + enc.to(self.device) + + # Move VAE to GPU if available + if modules.vae is not None: + try: + modules.vae.to(self.device, non_blocking=True) + except Exception as exc: + logger.debug("Failed to move VAE to GPU: %s", exc) + + logger.info("Applying layer-wise offloading on %s", modules.dit_names) + + # Apply block-wise offloading hook for each of the blocks in DiT model(s) + # Note that there might exist multiple DiT models in specific pipelines + for i, dit_module in enumerate(modules.dits): + dit_name = modules.dit_names[i] + logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})") + + blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(dit_module) + blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module) + + if not blocks_attr_name or not blocks: + logger.warning( + "Target layers (blocks) not found. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + num_blocks = len(blocks) + if num_blocks <= 1: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + # Move non-block modules to GPU (they stay resident) + for name, m in dit_module.named_children(): + if name == blocks_attr_name: + logger.debug(f"Skipped blocks module {name}") + continue + m.to(self.device) + logger.debug(f"Moved {name} to device {self.device}") + + # Pre-fetch the first layer by manually calling the hook function on the last layer; + # For subsequent requests, the first layer/block will be pre-fetched + # during the last layer compute of the previous request. + last_block, first_block = blocks[-1], blocks[0] + hook = apply_block_hook(last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory) + hook.prefetch_layer(non_blocking=False) + + # Register hook for each of blocks + for i, block in enumerate(blocks[:-1]): + next_block = blocks[(i + 1) % num_blocks] + apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory) + + logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)") + + # Track hooked blocks for cleanup + self._blocks.append(blocks) + + if len(self._blocks) > 0 and len(self._blocks[0]) > 0: + self.enabled = True + + def disable(self) -> None: + if not self.enabled: + return + + for blocks in self._blocks: + for block in blocks: + remove_block_hook(block) + + self._blocks.clear() + self.enabled = False + logger.info("Layer-wise offloading disabled") + + @staticmethod + def get_blocks_attr_name(model: nn.Module) -> str | None: + """Retrieve blocks attribute name from provided DiT model""" + return getattr(model.__class__, "_layerwise_offload_blocks_attr", None) + + @staticmethod + def set_blocks_attr_name(model: nn.Module, name: str) -> None: + if not hasattr(model.__class__, "_layerwise_offload_blocks_attr"): + setattr(model.__class__, "_layerwise_offload_blocks_attr", name) + + @staticmethod + def get_blocks_from_dit(model: nn.Module) -> list[nn.Module]: + """ + Retrieve a list of blocks from provided DiT model. Blocks attribute name + are found by `_layerwise_offload_blocks_attr` set to DiT models. For example, + + ``` + class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attr = "blocks" + ``` + """ + blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(model) + if blocks_attr_name is None: + logger.warning( + f"No _layerwise_offload_blocks_attr defined for {model.__class__.__name__}, " + "skipping layerwise offloading" + ) + return [] + + _blocks = getattr(model, blocks_attr_name, None) + if _blocks is None: + logger.warning( + f"Blocks (layers) '{blocks_attr_name}' not found on {model.__class__.__name__}, " + "skipping layerwise offloading" + ) + return [] + + return list(_blocks) diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py new file mode 100644 index 000000000..307ca53a8 --- /dev/null +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + +from torch import nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class PipelineModules: + dits: list[nn.Module] + dit_names: list[str] + encoders: list[nn.Module] + encoder_names: list[str] + vae: nn.Module | None = None + + +class ModuleDiscovery: + """Discovers pipeline components for offloading""" + + DIT_ATTRS = ["transformer", "transformer_2", "dit", "language_model", "transformer_blocks"] + ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"] + VAE_ATTRS = ["vae"] + + @staticmethod + def discover(pipeline: nn.Module) -> PipelineModules: + """Discover DiT, encoder, and VAE modules from pipeline. + + Args: + pipeline: Diffusion pipeline model + + Returns: + PipelineModules with lists of discovered modules and names + """ + # Collect DiT/transformer modules + dit_modules: list[nn.Module] = [] + dit_names: list[str] = [] + for attr in ModuleDiscovery.DIT_ATTRS: + if not hasattr(pipeline, attr): + continue + module_obj = getattr(pipeline, attr) + if module_obj is None: + continue + + if not isinstance(module_obj, nn.Module): + logger.warning(f"Expected {attr} to be nn.Module, got {type(module_obj)!r}") + continue + + if module_obj in dit_modules: + continue + + dit_modules.append(module_obj) + dit_names.append(attr) + + # Collect all encoders + encoders: list[nn.Module] = [] + encoder_names: list[str] = [] + for attr in ModuleDiscovery.ENCODER_ATTRS: + if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None: + encoders.append(getattr(pipeline, attr)) + encoder_names.append(attr) + + # Collect VAE + vae = None + for attr in ModuleDiscovery.VAE_ATTRS: + module = getattr(pipeline, attr, None) + if module is not None: + vae = module + break + + return PipelineModules( + dits=dit_modules, + dit_names=dit_names, + encoders=encoders, + encoder_names=encoder_names, + vae=vae, + ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py new file mode 100644 index 000000000..30edc8256 --- /dev/null +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.platforms import current_omni_platform + +from .base import OffloadBackend, OffloadConfig +from .module_collector import ModuleDiscovery + +logger = init_logger(__name__) + + +class SequentialOffloadHook(ModelHook): + """Hook for sequential offloading with mutual exclusion on encoder and DiT modules. + + To be used as a model-level (or "component-level") of CPU offloading method; + When a module's forward is called, this hook offloads target modules to CPU + and loads the current module to GPU. + """ + + _HOOK_NAME = "sequential_offload" + + def __init__( + self, + offload_targets: list[nn.Module], + device: torch.device, + pin_memory: bool = True, + ): + # Modules to offload to CPU before this module runs + self.offload_targets = offload_targets + self.device = device + self.pin_memory = pin_memory + + def _to_cpu(self, module: nn.Module) -> None: + """Move module to CPU.""" + try: + param = next(module.parameters()) + except StopIteration: + return + + previous_device = param.device + # Skip if already on CPU + if previous_device.type == "cpu": + return + + module.to("cpu") + torch.cuda.empty_cache() + + if self.pin_memory: + for p in module.parameters(): + if p.data.device.type == "cpu" and not p.data.is_pinned(): + p.data = p.data.pin_memory() + + def _to_gpu(self, module: nn.Module) -> None: + """Move module to GPU.""" + try: + # Skip if already on target device + if next(module.parameters()).device == self.device: + return + except StopIteration: + return + + module.to(self.device) + + def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]: + # Offload target modules to CPU + for target in self.offload_targets: + self._to_cpu(target) + + # Load current module to GPU + self._to_gpu(module) + current_omni_platform.synchronize() + + logger.debug( + "Swapped: %s -> CPU, %s -> GPU", + [t.__class__.__name__ for t in self.offload_targets], + module.__class__.__name__, + ) + + return args, kwargs + + +def apply_sequential_offload( + dit_modules: list[nn.Module], + encoder_modules: list[nn.Module], + device: torch.device, + pin_memory: bool = True, +) -> None: + """Apply sequential offloading hooks to DiT and encoder modules. + + Registers hooks on modules to implement mutual-exclusion GPU allocation. + - Before DiT runs, encoders are offloaded to CPU. + - Before encoders run, DiT is offloaded to CPU. + + Args: + dit_modules: DiT/transformer modules to register hooks on + encoder_modules: Encoder modules to register hooks on + device: Target GPU device for loading + pin_memory: Whether to pin CPU memory for faster transfers + + Example: + >>> apply_sequential_offload( + ... dit_modules=[pipeline.transformer], + ... encoder_modules=[pipeline.text_encoder, pipeline.vae], + ... device=torch.device("cuda:0"), + ... ) + >>> # Modules of pipeline now automatically swap between CPU and GPU + """ + # Register hooks on DiT modules (offload encoders when DiT runs) + for dit_mod in dit_modules: + registry = HookRegistry.get_or_create(dit_mod) + hook = SequentialOffloadHook( + offload_targets=encoder_modules, + device=device, + pin_memory=pin_memory, + ) + registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) + logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) + + # Register hooks on encoders (offload DiTs when encoder runs) + for enc in encoder_modules: + registry = HookRegistry.get_or_create(enc) + hook = SequentialOffloadHook( + offload_targets=dit_modules, + device=device, + pin_memory=pin_memory, + ) + registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) + logger.debug("Registered offload hook for %s", enc.__class__.__name__) + + +def remove_sequential_offload(modules: list[nn.Module]) -> None: + """Remove sequential offloading hooks from modules. + + Args: + modules: Modules to remove hooks from + + Example: + >>> all_modules = [*dit_modules, *encoder_modules] + >>> remove_sequential_offload(all_modules) + """ + for module in modules: + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is not None: + registry.remove_hook(SequentialOffloadHook._HOOK_NAME) + logger.debug("Removed offload hook from %s", module.__class__.__name__) + + +class ModelLevelOffloadBackend(OffloadBackend): + """Model-level (sequential) offloading backend. + + Uses SequentialOffloadHook registered via HookRegistry for automatic module swapping. + """ + + def __init__(self, config: OffloadConfig, device: torch.device): + super().__init__(config, device) + self._offload_modules: list[nn.Module] = [] # Track modules with hooks + + def enable(self, pipeline: nn.Module) -> None: + if self.enabled: + logger.warning("ModelLevelOffloadBackend already enabled") + return + + modules = ModuleDiscovery.discover(pipeline) + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping model-level offloading") + return + if not modules.encoders: + logger.warning("No encoder modules found, skipping model-level offloading") + return + + # Move encoders to GPU + for enc in modules.encoders: + enc.to(self.device) + + # Move VAE to GPU if available + if modules.vae is not None: + try: + modules.vae.to(self.device, non_blocking=True) + except Exception as exc: + logger.debug("Failed to move VAE to GPU: %s", exc) + + # Initial state: keep DiT modules on CPU (encoders typically run first) + # TODO: This part seems to be unnecessary, remove it after testing + # for dit_mod in modules.dits: + # dit_mod.to("cpu") + + # torch.cuda.empty_cache() + + # if self.config.pin_cpu_memory: + # for dit_mod in modules.dits: + # for p in dit_mod.parameters(): + # if p.data.device.type == "cpu" and not p.data.is_pinned(): + # p.data = p.data.pin_memory() + + # Apply sequential offloading hooks + apply_sequential_offload( + dit_modules=modules.dits, + encoder_modules=modules.encoders, + device=self.device, + pin_memory=self.config.pin_cpu_memory, + ) + + # Track modules for cleanup + self._offload_modules = [*modules.dits, *modules.encoders] + + self.enabled = True + + logger.info( + "Model-level offloading enabled: %s <-> %s (mutual exclusion)", + ", ".join(modules.dit_names), + ", ".join(modules.encoder_names), + ) + + def disable(self) -> None: + if not self.enabled: + return + + remove_sequential_offload(self._offload_modules) + + self._offload_modules.clear() + self.enabled = False + logger.info("Model-level offloading disabled") diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 1e00e4b0c..134a2ae39 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -26,7 +26,7 @@ from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.offload import apply_offload_hooks +from vllm_omni.diffusion.offloader import get_offload_backend from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.platforms import current_omni_platform @@ -62,6 +62,7 @@ class DiffusionModelRunner: self.device = device self.pipeline = None self.cache_backend = None + self.offload_backend = None # Initialize KV cache manager for connector management self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config) @@ -69,6 +70,8 @@ class DiffusionModelRunner: def load_model( self, memory_pool_context_fn: callable | None = None, + load_format: str | None = None, + custom_pipeline_name: str | None = None, ) -> None: """ Load the diffusion model, apply compilation and offloading. @@ -76,7 +79,17 @@ class DiffusionModelRunner: Args: memory_pool_context_fn: Optional function that returns a context manager for memory pool allocation (used for sleep mode). + load_format: Format for loading model weights. Supported formats: + - "default" (default): Automatically detect and use the default format based on configuration + - "custom_pipeline": Init model from a custom pipeline class specified by `custom_pipeline_name` + - "dummy": Skip actual weight loading, useful for testing and custom pipelines that + don't require default weights. + custom_pipeline_name: Optional custom pipeline class name to use. """ + + if load_format == "dummy": + return + load_device = ( "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device) ) @@ -96,6 +109,8 @@ class DiffusionModelRunner: self.pipeline = model_loader.load_model( od_config=self.od_config, load_device=load_device, + load_format=load_format, + custom_pipeline_name=custom_pipeline_name, ) time_after_load = time.perf_counter() @@ -107,8 +122,10 @@ class DiffusionModelRunner: logger.info("Model runner: Model loaded successfully.") # Apply CPU offloading - if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload: - apply_offload_hooks(self.pipeline, self.od_config, device=self.device) + self.offload_backend = get_offload_backend(self.od_config, device=self.device) + if self.offload_backend is not None: + logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}") + self.offload_backend.enable(self.pipeline) # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 6f5bd00e6..6a8494991 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """ Diffusion Worker for vLLM-Omni. @@ -7,15 +8,19 @@ Handles GPU infrastructure initialization and delegates model operations to DiffusionModelRunner. """ +import gc import multiprocessing as mp import os +from collections.abc import Iterable from contextlib import AbstractContextManager, nullcontext +from typing import Any import torch import zmq from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.mem_utils import GiB_bytes from vllm.v1.worker.workspace import init_workspace_manager @@ -35,6 +40,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.lora.request import LoRARequest from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory logger = init_logger(__name__) @@ -67,6 +73,15 @@ class DiffusionWorker: self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self.lora_manager: DiffusionLoRAManager | None = None self.init_device() + # Create model runner + self.model_runner = DiffusionModelRunner( + vllm_config=self.vllm_config, + od_config=self.od_config, + device=self.device, + ) + self.load_model(load_format=self.od_config.diffusion_load_format) + self.init_lora_manager() + logger.info(f"Worker {self.rank}: Initialization complete.") def init_device(self) -> None: """Initialize the device and distributed environment.""" @@ -110,16 +125,30 @@ class DiffusionWorker: ) init_workspace_manager(self.device) - # Create model runner and load model - self.model_runner = DiffusionModelRunner( - vllm_config=self.vllm_config, - od_config=self.od_config, - device=self.device, - ) + def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None: + """Load the diffusion model using DiffusionModelRunner.""" + with ( + set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config), + set_current_vllm_config(self.vllm_config), + ): self.model_runner.load_model( memory_pool_context_fn=self._maybe_get_memory_pool_context, + load_format=load_format, + custom_pipeline_name=custom_pipeline_name, + ) + process_memory = get_process_gpu_memory(self.local_rank) + if process_memory is not None: + logger.info( + "Worker %d: Process-scoped GPU memory after model loading: %.2f GiB.", + self.rank, + process_memory / GiB_bytes, ) assert self.model_runner.pipeline is not None + + def init_lora_manager(self) -> None: + """Initialize the LoRA manager for this worker.""" + if self.model_runner.pipeline is None: + return self.lora_manager = DiffusionLoRAManager( pipeline=self.model_runner.pipeline, device=self.device, @@ -128,7 +157,6 @@ class DiffusionWorker: lora_path=self.od_config.lora_path, lora_scale=self.od_config.lora_scale, ) - logger.info(f"Worker {self.rank}: Initialization complete.") def generate(self, request: OmniDiffusionRequest) -> DiffusionOutput: """Generate output for the given requests.""" @@ -182,7 +210,10 @@ class DiffusionWorker: """ from vllm.device_allocator.cumem import CuMemAllocator - free_bytes_before_sleep = current_omni_platform.get_free_memory() + process_memory_before_sleep = get_process_gpu_memory(self.local_rank) + free_bytes_before_sleep = None + if process_memory_before_sleep is None: + free_bytes_before_sleep = current_omni_platform.get_free_memory() # Save the buffers before level 2 sleep if level == 2 and self.model_runner is not None: @@ -191,14 +222,23 @@ class DiffusionWorker: allocator = CuMemAllocator.get_instance() allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) - free_bytes_after_sleep = current_omni_platform.get_free_memory() - device_id = self.device.index if self.device.index is not None else 0 - total = current_omni_platform.get_device_total_memory(device_id) - freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep - used_bytes = total - free_bytes_after_sleep + process_memory_after_sleep = get_process_gpu_memory(self.local_rank) + if process_memory_before_sleep is not None and process_memory_after_sleep is not None: + freed_bytes = process_memory_before_sleep - process_memory_after_sleep + used_bytes = process_memory_after_sleep + accounting_scope = "process-scoped" + else: + free_bytes_after_sleep = current_omni_platform.get_free_memory() + assert free_bytes_before_sleep is not None + device_id = self.device.index if self.device.index is not None else 0 + total = current_omni_platform.get_device_total_memory(device_id) + freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep + used_bytes = total - free_bytes_after_sleep + accounting_scope = "device-scoped fallback" assert freed_bytes >= 0, "Memory usage increased after sleeping." logger.info( - "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.", + "Sleep mode (%s) freed %.2f GiB memory, %.2f GiB memory is still in use.", + accounting_scope, freed_bytes / GiB_bytes, used_bytes / GiB_bytes, ) @@ -247,6 +287,32 @@ class DiffusionWorker: destroy_distributed_env() +class CustomPipelineWorkerExtension: + def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]) -> None: + """ + Re-initialize the pipeline with custom arguments. + + Args: + custom_pipeline_args: Dictionary of arguments for custom pipeline initialization + """ + + # Clean up old pipeline + if self.model_runner.pipeline is not None: + del self.model_runner.pipeline + gc.collect() + torch.cuda.empty_cache() + + # Get custom pipeline class name + custom_pipeline_name = custom_pipeline_args["pipeline_class"] + + # Use the DiffusionWorker's load_model method which handles the forward context + self.load_model( + load_format="custom_pipeline", + custom_pipeline_name=custom_pipeline_name, + ) + self.init_lora_manager() + + class WorkerProc: """Wrapper that runs one Worker in a separate process.""" @@ -255,8 +321,11 @@ class WorkerProc: od_config: OmniDiffusionConfig, gpu_id: int, broadcast_handle, + worker_extension_cls: str | None = None, + custom_pipeline_args: dict[str, Any] | None = None, ): self.od_config = od_config + self.gpu_id = gpu_id # Inter-process Communication self.context = zmq.Context(io_threads=2) @@ -274,17 +343,26 @@ class WorkerProc: logger.info(f"Worker {gpu_id} created result MessageQueue") assert od_config.master_port is not None - self.worker = self._create_worker(gpu_id, od_config) - self.gpu_id = gpu_id + + # Create worker using WorkerWrapperBase for extension support + self.worker = self._create_worker(gpu_id, od_config, worker_extension_cls, custom_pipeline_args) self._running = True - def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> DiffusionWorker: + def _create_worker( + self, + gpu_id: int, + od_config: OmniDiffusionConfig, + worker_extension_cls: str | None, + custom_pipeline_args: dict[str, Any] | None = None, + ) -> DiffusionWorker: """Create a worker instance. Override in subclasses for different worker types.""" - return DiffusionWorker( - local_rank=gpu_id, - rank=gpu_id, + wrapper = WorkerWrapperBase( + gpu_id=gpu_id, od_config=od_config, + worker_extension_cls=worker_extension_cls, + custom_pipeline_args=custom_pipeline_args, ) + return wrapper def return_result(self, output: DiffusionOutput): """Reply to client, only on rank 0.""" @@ -310,11 +388,8 @@ class WorkerProc: return None, False try: - if isinstance(method, str): - func = getattr(self.worker, method) - result = func(*args, **kwargs) - else: - result = method(self.worker, *args, **kwargs) + # Use execute_method from WorkerWrapperBase for consistent method resolution + result = self.worker.execute_method(method, *args, **kwargs) return result, should_reply except Exception as e: logger.error(f"Error executing RPC: {e}", exc_info=True) @@ -385,6 +460,8 @@ class WorkerProc: od_config: OmniDiffusionConfig, pipe_writer: mp.connection.Connection, broadcast_handle, + worker_extension_cls: str | None = None, + custom_pipeline_args: dict[str, Any] | None = None, ) -> None: """Worker initialization and execution loops.""" from vllm_omni.plugins import load_omni_general_plugins @@ -394,6 +471,8 @@ class WorkerProc: od_config, gpu_id=rank, broadcast_handle=broadcast_handle, + worker_extension_cls=worker_extension_cls, + custom_pipeline_args=custom_pipeline_args, ) logger.info(f"Worker {rank}: Scheduler loop started.") pipe_writer.send( @@ -404,3 +483,201 @@ class WorkerProc: ) worker_proc.worker_busy_loop() logger.info(f"Worker {rank}: Shutdown complete.") + + +class WorkerWrapperBase: + """ + Wrapper base class that creates DiffusionWorker with optional worker_extension_cls support. + This enables dynamic inheritance for DiffusionWorker to extend with custom functionality. + """ + + def __init__( + self, + gpu_id: int, + od_config: OmniDiffusionConfig, + base_worker_class: type = DiffusionWorker, + worker_extension_cls: str | None = None, + custom_pipeline_args: dict[str, Any] | None = None, + ): + """ + Initialize WorkerWrapperBase with support for worker extensions. + + Args: + gpu_id: GPU device ID + od_config: OmniDiffusionConfig configuration + worker_extension_cls: Optional qualified name of worker extension class + custom_pipeline_args: Optional arguments for custom pipeline initialization + """ + self.gpu_id = gpu_id + self.od_config = od_config + self.base_worker_class = base_worker_class + self.worker_extension_cls = worker_extension_cls + self.custom_pipeline_args = custom_pipeline_args + + # Prepare worker class with extension support + worker_class = self._prepare_worker_class() + + # Create the actual worker instance + self.worker = worker_class( + local_rank=gpu_id, + rank=gpu_id, + od_config=od_config, + ) + + # Re-initialize pipeline with custom pipeline if provided + if self.custom_pipeline_args is not None: + self.worker.re_init_pipeline(self.custom_pipeline_args) + + def _prepare_worker_class(self) -> type: + """ + Prepare the worker class with optional extension. + Dynamically extends GPUWorker with worker_extension_cls if provided. + + Returns: + The worker class (potentially extended) + """ + worker_class = self.base_worker_class + + # If custom_pipeline_args is provided, use CustomPipelineWorkerExtension + if self.custom_pipeline_args is not None: + # Set worker_extension_cls to CustomPipelineWorkerExtension if not already set + if self.worker_extension_cls is None: + self.worker_extension_cls = CustomPipelineWorkerExtension + + if self.worker_extension_cls: + if isinstance(self.worker_extension_cls, str): + worker_extension_cls = resolve_obj_by_qualname(self.worker_extension_cls) + else: + worker_extension_cls = self.worker_extension_cls + extended_calls = [] + + if worker_extension_cls not in worker_class.__bases__: + # Check for conflicts between worker and extension + for attr in dir(worker_extension_cls): + if attr.startswith("__"): + continue + if hasattr(worker_class, attr): + logger.warning( + f"Worker class {worker_class} already has attribute " + f"{attr}, which may conflict with worker extension " + f"class {worker_extension_cls}." + ) + if callable(getattr(worker_extension_cls, attr)): + extended_calls.append(attr) + + # Dynamically inherit the worker extension class + class_name = f"{worker_class.__name__}With{worker_extension_cls.__name__}" + worker_class = type(class_name, (worker_extension_cls, worker_class), {}) + logger.info( + "Created extended worker class %s from %s for extended calls %s", + class_name, + worker_extension_cls, + extended_calls, + ) + + return worker_class + + def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput: + """ + Generate output for the given requests. + + Args: + requests: List of diffusion requests + + Returns: + DiffusionOutput with generated results + """ + return self.worker.generate(requests) + + def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput: + """ + Execute a forward pass. + + Args: + reqs: List of diffusion requests + od_config: OmniDiffusionConfig configuration + + Returns: + DiffusionOutput with generated results + """ + return self.worker.execute_model(reqs, od_config) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """ + Load model weights. + + Args: + weights: Iterable of (name, tensor) tuples + + Returns: + Set of loaded weight names + """ + return self.worker.load_weights(weights) + + def sleep(self, level: int = 1) -> bool: + """ + Put the worker to sleep. The worker should not process any requests. + The caller should guarantee that no requests are being processed + during the sleep period, before `wake_up` is called. + + Args: + level: The sleep level. Level 1 sleep will offload the model + weights and discard the kv cache. + Currently only support level 1. + + Returns: + True on success + """ + return self.worker.sleep(level) + + def wake_up(self, tags: list[str] | None = None) -> bool: + """ + Wake up the worker from sleep mode. See the sleep function + method for more details. + + Args: + tags: An optional list of tags to reallocate the worker memory + for specific memory allocations. Values must be in + `("weights")`. If None, all memory is reallocated. + wake_up should be called with all tags (or None) before the + worker is used again. + + Returns: + True on success + """ + return self.worker.wake_up(tags) + + def shutdown(self) -> None: + """Shutdown the worker and cleanup resources.""" + return self.worker.shutdown() + + def execute_method(self, method: str | bytes, *args, **kwargs) -> Any: + """ + Execute a method on the worker. + + Args: + method: Method name (str) or serialized callable (bytes) + + Returns: + Result of the method execution (type depends on the method) + + Raises: + Exception: If method execution fails + """ + try: + # Method resolution order: + # 1. If method is defined in this class, it will be called directly + # 2. Otherwise, since we define `__getattr__` and redirect attribute + # query to `self.worker`, the method will be called on the worker + assert isinstance(method, str), "Method must be str" + func = getattr(self.worker, method) + return func(*args, **kwargs) + + except Exception as e: + msg = f"Error executing method {method!r}. This might cause issues in distributed execution." + logger.exception(msg) + raise e + + def __getattr__(self, attr: str): + """Delegate attribute access to the wrapped worker.""" + return getattr(self.worker, attr) diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 66aebddc1..1d1dd6f0f 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -5,14 +5,7 @@ import time from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -import torch -from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.request import Request, RequestStatus - -if TYPE_CHECKING: - from .connectors.base import OmniConnectorBase +from typing import Any from vllm_omni.entrypoints.stage_utils import OmniStageTaskType from vllm_omni.metrics import OrchestratorAggregator @@ -180,178 +173,6 @@ def try_recv_via_connector( return None, None -def update_request_payload(connector: "OmniConnectorBase", req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]: - """Update the payload data for a request in the connector. - - Args: - connector: OmniConnectorBase instance - req_id: Request ID to update - payload_data: New payload data to store - """ - origin_payload = connector.request_payload[req_id] - for key, value in payload_data.items(): - if key == "finished": - continue - elif isinstance(value, torch.Tensor) and key in origin_payload: - payload_data[key] = torch.cat([origin_payload[key], value], dim=0) - elif isinstance(value, list) and key in origin_payload: - payload_data[key] = origin_payload[key] + value - - connector.request_payload[req_id] = payload_data - return payload_data - - -def get_chunk( - connector: "OmniConnectorBase", - scheduler_output: SchedulerOutput, -) -> None: - """Retrieve a chunk of pooling output. - - Args: - connector: OmniConnectorBase instance - scheduler_output: Partial scheduler output dictionary - - Returns: - None: This function modifies scheduler_output in place - """ - stage_id = connector.stage_id - if stage_id == 0: - return - - target_stage_id = stage_id - 1 - # Handle new requests - for new_req_data in scheduler_output.scheduled_new_reqs: - connector.request_ids_mapping[new_req_data.req_id] = new_req_data.external_req_id - req_id = new_req_data.external_req_id - chunk_id = connector.get_requests[req_id] - connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" - payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) - if payload_data: - new_req_data.additional_information = payload_data - connector.request_payload[req_id] = payload_data - if payload_data.get("finished"): - connector.finished_requests.add(req_id) - - # Handle cached/running requests - cached_reqs = scheduler_output.scheduled_cached_reqs - if not hasattr(cached_reqs, "additional_information"): - cached_reqs.additional_information = {} - - for i, cached_req_id in enumerate(cached_reqs.req_ids): - req_id = connector.request_ids_mapping.get(cached_req_id, cached_req_id) - if req_id in connector.finished_requests: - continue - chunk_id = connector.get_requests[req_id] - connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}" - payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key) - if payload_data: - payload_data = update_request_payload(connector, req_id, payload_data) - cached_reqs.additional_information[cached_req_id] = payload_data - if payload_data.get("finished"): - connector.finished_requests.add(req_id) - - -def get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key): - # TODO: add correct check mechanism for the payload_data - max_wait = 300 - for _ in range(max_wait): - result = connector.get( - from_stage=str(target_stage_id), - to_stage=str(stage_id), - get_key=connector_get_key, - ) - payload_data = None - if result: - payload_data, size = result - if payload_data: - connector.request_prompt_token_ids[req_id] = payload_data.get("thinker_input_ids", []) - connector.get_requests[req_id] += 1 - logger.debug("[Stage-%d] Received one chunk for request %s", stage_id, connector_get_key) - break - time.sleep(0.01) - return payload_data - - -def get_chunk_for_generation( - connector: "OmniConnectorBase", - request: Request, -) -> None: - """Retrieve a chunk of pooling output. - - Args: - connector: OmniConnectorBase instance - request: Request object - - Returns: - None: This function modifies request in place - """ - stage_id = connector.stage_id - target_stage_id = stage_id - 1 - request_id = request.external_req_id - - if request_id in connector.finished_requests: - return - - chunk_id = connector.get_requests[request_id] - connector_get_key = f"{request_id}_{target_stage_id}_{chunk_id}" - payload_data = get_through_connector(connector, target_stage_id, stage_id, request_id, connector_get_key) - if not payload_data: - return - - if payload_data.get("finished"): - connector.finished_requests.add(request_id) - request.status = RequestStatus.FINISHED_STOPPED - request.prompt_token_ids = payload_data.get("code_predictor_codes", []) - request.num_computed_tokens = 0 - - -def put_chunk( - connector: "OmniConnectorBase", - pooling_output: dict[str, Any], - request: Request, - custom_process_input_func: Callable[[dict[str, Any], Request], dict[str, Any] | None] | None = None, -) -> None: - """Store a chunk of pooling output. - - Args: - connector: OmniConnectorBase instance - pooling_output: Partial pooling output dictionary - request: Request object - custom_process_input_func: Optional custom function to process input - - Returns: - None: This function sends data via connector - """ - stage_id = connector.stage_id - next_stage_id = stage_id + 1 - request_id = request.external_req_id - chunk_id = connector.put_requests[request_id] - connector_put_key = f"{request_id}_{stage_id}_{chunk_id}" - payload_data = None - - # TODO: add default process_input_func to handle the payload_data ? - if custom_process_input_func: - try: - payload_data = custom_process_input_func( - connector=connector, - pooling_output=pooling_output, - request=request, - ) - except Exception as e: - logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}") - - if not payload_data: - return - - success, size, metadata = connector.put( - from_stage=str(stage_id), to_stage=str(next_stage_id), put_key=connector_put_key, data=payload_data - ) - - if success: - connector.put_requests[request_id] += 1 - logger.debug("[Stage-%d] Sent %s", stage_id, connector_put_key) - - def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: """Compute the length of the talker prompt ids. diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index a11f38832..c654e2697 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -4,7 +4,6 @@ import fcntl import os import time -from collections import defaultdict from typing import Any from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes @@ -25,14 +24,6 @@ class SharedMemoryConnector(OmniConnectorBase): self.config = config self.stage_id = config.get("stage_id", -1) self.device = config.get("device", "cuda:0") - self.put_requests: dict[str, int] = defaultdict(int) - self.get_requests: dict[str, int] = defaultdict(int) - self.finished_requests: set[str] = set() - self.request_payload = {} - self.request_prompt_token_ids: dict[str, list[int]] = defaultdict(list) - self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) - self.request_ids_mapping: dict[str, str] = {} - # Default threshold matches legacy behavior (64KB) self.threshold = int(config.get("shm_threshold_bytes", 65536)) self._metrics = { "puts": 0, diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/__init__.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py new file mode 100644 index 000000000..0962808ed --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +import time +from typing import Any + +from ..utils.logging import get_connector_logger + +logger = get_connector_logger(__name__) + + +class OmniTransferAdapterBase: + """Base class for managing data transfer via OmniConnector. + + This class handles the core loop logic and connector interactions, but + leaves the specific data processing (chunks, KV cache, etc.) to subclasses. + """ + + def __init__(self, config: Any): + self.config = config + if not hasattr(self, "connector"): + self.connector = None + # Requests that are waiting to be polled + self._pending_load_reqs = {} + # Requests that have successfully retrieved data + self._finished_load_reqs = set() + + # Requests that are waiting to be saved + self._pending_save_reqs = {} + # Requests that have successfully saved data + self._finished_save_reqs = set() + + self.stop_event = threading.Event() + self.lock = threading.Lock() + + self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True) + self.recv_thread.start() + + self.save_thread = threading.Thread(target=self.save_loop, daemon=True) + self.save_thread.start() + + @classmethod + def create_connector(cls, model_config: Any): + raise NotImplementedError + + def recv_loop(self): + """Loop to poll for incoming data.""" + while not self.stop_event.is_set(): + # Iterate over a snapshot of pending requests + with self.lock: + pending_reqs_ids = list(self._pending_load_reqs.keys()) + + for req_id in pending_reqs_ids: + try: + self._poll_single_request(req_id) + except Exception as e: + logger.warning(f"Error receiving data for {req_id}: {e}") + + time.sleep(0.001) + + def save_loop(self): + """Loop to send outgoing data.""" + while not self.stop_event.is_set(): + task = None + with self.lock: + pending_save_reqs_ids = list(self._pending_save_reqs.keys()) + for req_id in pending_save_reqs_ids: + if self._pending_save_reqs[req_id]: + task = self._pending_save_reqs[req_id].popleft() + if not self._pending_save_reqs[req_id]: + del self._pending_save_reqs[req_id] + break + + if task: + try: + self._send_single_request(task) + except Exception as e: + logger.error(f"Error saving data for {task.get('request_id')}: {e}") + else: + time.sleep(0.001) + + def _poll_single_request(self, *args, **kwargs): + """Poll connector for a single request task. + Subclasses should implement request-specific receive behavior.""" + raise NotImplementedError + + def _send_single_request(self, *args, **kwargs): + """Send one pending save request task to the connector. + Subclasses should implement task-specific handling logic.""" + raise NotImplementedError + + def load_async(self, *args, **kwargs): + """Register a request to load data. To be implemented by subclasses.""" + raise NotImplementedError + + def save_async(self, *args, **kwargs): + """Submit data to be saved. To be implemented by subclasses.""" + raise NotImplementedError + + def load(self, *args, **kwargs): + """Load request data from connector synchronously. To be implemented by subclasses.""" + raise NotImplementedError + + def save(self, *args, **kwargs): + """Save data to connector synchronously. To be implemented by subclasses.""" + raise NotImplementedError + + def get_finished_requests(self): + """Get finished loaded or saved requests""" + raise NotImplementedError + + def shutdown(self): + """Stop background loops and close the connector.""" + raise NotImplementedError diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py new file mode 100644 index 000000000..70d38a9d6 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from collections import defaultdict, deque +from typing import Any + +import torch +from vllm.v1.request import Request, RequestStatus + +from ..factory import OmniConnectorFactory +from ..utils.config import ConnectorSpec +from ..utils.logging import get_connector_logger +from .base import OmniTransferAdapterBase + +logger = get_connector_logger(__name__) + + +class OmniChunkTransferAdapter(OmniTransferAdapterBase): + """Chunk-level transfer adapter for Omni connector pipelines. + + This class coordinates per-request chunk exchange between adjacent stages, + and implements asynchronous get/put of chunks via background threads. + It tracks per-request chunk indices for put/get, and accumulates + payloads across chunks (concatenating tensors/lists in AR mode). It also + caches prompt token ids and additional information for scheduler use. + + Scheduler integration is handled via WAITING_FOR_CHUNK transitions: + requests are moved to waiting for chunk deque while polling, then restored + to waiting/running queues once a chunk arrives. The requests will finish + loading chunk util detecting the payload "finished" flag. + + The base class owns background recv/save loops; load/save only enqueue + work and return immediately. + """ + + def __init__(self, vllm_config: Any): + model_config = vllm_config.model_config + self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs + self.connector = self.create_connector(model_config) + super().__init__(model_config) + self.model_mode = getattr(model_config, "worker_type", "ar") + # State specific to Chunk management + self.custom_process_next_stage_input_func = None + custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) + if custom_process_next_stage_input_func: + module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) + module = importlib.import_module(module_path) + self.custom_process_next_stage_input_func = getattr(module, func_name) + # mapping for request id and chunk id + self.put_req_chunk: dict[str, int] = defaultdict(int) + self.get_req_chunk: dict[str, int] = defaultdict(int) + self.finished_requests: set[str] = set() + self.request_payload = {} + self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) + self.request_ids_mapping: dict[str, str] = {} + + self.waiting_for_chunk_waiting_requests: deque[Any] = deque() + self.waiting_for_chunk_running_requests: deque[Any] = deque() + self.requests_with_ready_chunks = set() + + @classmethod + def create_connector(cls, model_config: Any): + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is None: + connector_config = {} + elif not isinstance(connector_config, dict): + connector_config = { + "name": getattr(connector_config, "name", None), + "extra": getattr(connector_config, "extra", {}), + } + + connector_specs = ConnectorSpec( + name=connector_config.get("name", "SharedMemoryConnector"), + extra=connector_config.get("extra", {}), + ) + return OmniConnectorFactory.create_connector(connector_specs) + + def load_async(self, request: Request): + """Register a request for asynchronous chunk retrieval. + + This method does not read from the connector directly. It records + request metadata and enqueues the request id for the background + receive loop to poll. + + Stage-0 has no upstream producer, so this call is a no-op there. + + Args: + request: The request object needing data. + """ + stage_id = self.connector.stage_id + request_id = request.request_id + self.request_ids_mapping[request_id] = request.external_req_id + + if stage_id == 0: + return + if not hasattr(request, "additional_information"): + request.additional_information = None + with self.lock: + self._pending_load_reqs[request_id] = request + + def save_async(self, pooling_output: torch.Tensor | None = None, request: Request | None = None): + """Build and enqueue one chunk for asynchronous sending. + + Payload extraction is executed in the caller thread via + ``custom_process_next_stage_input_func`` + + Args: + pooling_output: Partial pooling output dictionary + request: Request object + """ + stage_id = self.connector.stage_id + next_stage_id = stage_id + 1 + request_id = request.external_req_id + chunk_id = self.put_req_chunk[request_id] + + # Process payload in main thread to avoid race conditions on request state + payload_data = None + if self.custom_process_next_stage_input_func: + try: + payload_data = self.custom_process_next_stage_input_func( + transfer_manager=self, + pooling_output=pooling_output, + request=request, + ) + + except Exception as e: + logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}") + + if not payload_data: + return + + # Increment chunk_id + self.put_req_chunk[request_id] += 1 + connector_put_key = f"{request_id}_{stage_id}_{chunk_id}" + + task = { + "stage_id": stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload_data, + "request_id": request_id, + } + + with self.lock: + if request_id not in self._pending_save_reqs: + self._pending_save_reqs[request_id] = deque() + self._pending_save_reqs[request_id].append(task) + + def _poll_single_request(self, req_id: str): + stage_id = self.connector.stage_id + target_stage_id = stage_id - 1 + chunk_id = self.get_req_chunk[req_id] + external_req_id = self.request_ids_mapping.get(req_id, req_id) + connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" + + # Use timeout=0 for non-blocking poll + result = self.connector.get( + str(target_stage_id), + str(stage_id), + connector_get_key, + ) + + if result is None: + return + + payload_data, size = result + + if payload_data: + # Update connector state + self.get_req_chunk[req_id] += 1 + req = self._pending_load_reqs[req_id] + + if self.model_mode == "ar": + self._update_request_payload(external_req_id, payload_data) + req.additional_information = payload_data + if payload_data.get("finished"): + self.finished_requests.add(req_id) + else: + if payload_data.get("finished"): + self.finished_requests.add(req_id) + req.status = RequestStatus.FINISHED_STOPPED + + req.prompt_token_ids = payload_data.get("code_predictor_codes", []) + req.num_computed_tokens = 0 + + # Mark as finished for consumption + with self.lock: + self._finished_load_reqs.add(req_id) + if req_id in self._pending_load_reqs: + del self._pending_load_reqs[req_id] + logger.debug(f"[Stage-{stage_id}] Received one chunk for key {connector_get_key}") + + def _update_request_payload(self, req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]: + """Update the payload data for a request in the connector. + + Args: + connector: OmniConnectorBase instance + req_id: Request ID to update + payload_data: New payload data to store + """ + if req_id not in self.request_payload: + self.request_payload[req_id] = payload_data + return payload_data + origin_payload = self.request_payload[req_id] + for key, value in payload_data.items(): + if key == "finished": + continue + elif isinstance(value, torch.Tensor) and key in origin_payload: + payload_data[key] = torch.cat([origin_payload[key], value], dim=0) + elif isinstance(value, list) and key in origin_payload: + payload_data[key] = origin_payload[key] + value + + self.request_payload[req_id] = payload_data + return payload_data + + def _send_single_request(self, task: dict): + connector_put_key = task["put_key"] + stage_id = task["stage_id"] + next_stage_id = task["next_stage_id"] + payload_data = task["data"] + + success, size, metadata = self.connector.put( + from_stage=str(stage_id), + to_stage=str(next_stage_id), + put_key=connector_put_key, + data=payload_data, + ) + + if success: + logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") + + ######################################################################## + # Schedule Helper + ######################################################################## + + def process_pending_chunks( + self, + waiting_queue: Any, + running_queue: list[Request], + ) -> None: + """ + Process pending chunks for waiting and running queues. + """ + if self.connector.stage_id == 0: + return + finished_load_reqs = self.get_finished_requests() + self._process_chunk_queue( + waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, finished_load_reqs + ) + self._process_chunk_queue( + running_queue, self.waiting_for_chunk_running_requests, RequestStatus.RUNNING, finished_load_reqs + ) + while len(running_queue) > self.scheduler_max_num_seqs: + request = running_queue.pop() + waiting_queue.prepend_requests([request]) + + def restore_queues(self, waiting_queue: Any, running_queue: list[Request]) -> None: + """ + Restore requests waiting for chunk to the waiting and running queues. + """ + # Add request waiting for chunk to the waiting and running queue + for request in self.waiting_for_chunk_waiting_requests: + waiting_queue.add_request(request) + self.waiting_for_chunk_waiting_requests = deque() + + if self.waiting_for_chunk_running_requests: + running_queue.extend(self.waiting_for_chunk_running_requests) + self.waiting_for_chunk_running_requests = deque() + + def postprocess_scheduler_output( + self, + scheduler_output: Any, + requests: dict[str, Request] | None = None, + ) -> None: + """ + Add additional info for cached requests and + clean up ready chunks from scheduler output. + """ + if requests is not None: + self.attach_cached_additional_information(scheduler_output, requests) + self._clear_chunk_ready(scheduler_output) + + @staticmethod + def attach_cached_additional_information(scheduler_output: Any, requests: dict[str, Request]) -> None: + cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None) + if not cached_reqs: + return + if not hasattr(cached_reqs, "additional_information"): + cached_reqs.additional_information = {} + for req_id in cached_reqs.req_ids: + request = requests.get(req_id) if req_id else None + additional_info = getattr(request, "additional_information", None) if request else None + cached_reqs.additional_information[req_id] = additional_info + + def _process_chunk_queue( + self, + queue: Any, + waiting_for_chunk_list: deque[Any], + target_status: RequestStatus, + finished_load_reqs: set[str], + ) -> None: + queue_snapshot = list(queue) + for request in queue_snapshot: + if request.status != RequestStatus.WAITING_FOR_CHUNK: + if request.request_id in self.requests_with_ready_chunks: + # Requests that have loaded chunk from last round + # of schedule, but have not scheduled + continue + if request.request_id in self.finished_requests: + request.additional_information = {} + continue + # Requests that waiting for chunk + self.load_async(request) + request.status = RequestStatus.WAITING_FOR_CHUNK + else: + if request.request_id in finished_load_reqs: + request.status = target_status + self.requests_with_ready_chunks.add(request.request_id) + continue + queue.remove(request) + waiting_for_chunk_list.append(request) + + def _clear_chunk_ready(self, scheduler_output: Any) -> None: + if scheduler_output.scheduled_new_reqs: + for req_data in scheduler_output.scheduled_new_reqs: + if req_data.req_id in self.requests_with_ready_chunks: + self.requests_with_ready_chunks.remove(req_data.req_id) + + if scheduler_output.scheduled_cached_reqs: + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + if req_id in self.requests_with_ready_chunks: + self.requests_with_ready_chunks.remove(req_id) + + def get_finished_requests(self): + with self.lock: + finished_load = set(self._finished_load_reqs) + self._finished_load_reqs = set() + return finished_load diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 4e1cb7c8c..f2adca975 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -62,6 +62,7 @@ class OmniEngineArgs(EngineArgs): If None, default processing is used. 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" """ stage_id: int = 0 @@ -73,6 +74,7 @@ class OmniEngineArgs(EngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) async_chunk: bool = False omni_kv_config: dict | None = None + worker_type: str | None = None def __post_init__(self) -> None: load_omni_general_plugins() @@ -217,6 +219,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. stage_connector_spec: Extra configuration for stage connector + worker_type: Model Type, e.g., "ar" or "generation" """ stage_id: int = 0 @@ -228,6 +231,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) async_chunk: bool = False omni_kv_config: dict | None = None + worker_type: str | None = None def __post_init__(self) -> None: load_omni_general_plugins() diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ad55c048d..369713d7b 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -162,8 +162,9 @@ class AsyncOmni(OmniBase): "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), - "layerwise_num_gpu_layers": kwargs.get("layerwise_num_gpu_layers", False), "enforce_eager": kwargs.get("enforce_eager", False), + "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), + "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), }, "final_output": True, "final_output_type": "image", diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 1be3b02f1..a59cc5673 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -204,12 +204,6 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - omni_config_group.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of layers (blocks) to keep on GPU during generation.", - ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 3116e3294..6c9723b6b 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -67,16 +67,20 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = """ from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available - if is_process_scoped_memory_available() and detect_pid_host(): - logger.debug( + nvml_available = is_process_scoped_memory_available() + pid_host = detect_pid_host() + + if nvml_available and pid_host: + logger.info( "NVML process-scoped memory available and PID host is available — concurrent init is safe, skipping locks" ) yield return else: - logger.debug( - "NVML unavailable or PID host is not available (usually inside a container, " - "--pid=host is not set in docker run command) — using sequential init locks" + logger.info( + "Using sequential init locks (nvml_available=%s, pid_host=%s)", + nvml_available, + pid_host, ) from vllm_omni.platforms import current_omni_platform @@ -201,7 +205,7 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = def _resolve_worker_cls(engine_args: dict[str, Any]) -> None: - worker_type = engine_args.pop("worker_type", None) + worker_type = engine_args.get("worker_type", None) if not worker_type: return if engine_args.get("worker_cls"): diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 0abe029d6..58d44919d 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -306,37 +306,42 @@ def in_container() -> bool: return any(m in cg for m in markers) -def has_pid_host() -> bool | None: - """ +def _nvml_pid_matches() -> bool | None: + """Check if os.getpid() matches what NVML would report via NSpid. + + If /proc/self/status has a single NSpid entry, the container PID equals + the host PID, so NVML process-scoped memory tracking works correctly. + Returns: - True -> very likely running with --pid=host (host PID namespace) - False -> very likely isolated PID namespace (default) + True -> PID namespaces match (single NSpid entry) + False -> PID namespaces differ (multiple NSpid entries) None -> cannot determine """ - # Strong signal: in host pid namespace, PID 2 is usually kthreadd - comm2 = _read_text("/proc/2/comm") - if comm2 is not None: - comm2 = comm2.strip() - if comm2 == "kthreadd": - return True - # If PID 2 exists and is NOT kthreadd, we're almost certainly not in host pid ns - return False - - # Fallback: check for other low-numbered kernel threads (best-effort) - for pid, name in [(3, "rcu_gp"), (4, "rcu_par_gp"), (10, "ksoftirqd/0")]: - comm = _read_text(f"/proc/{pid}/comm") - if comm is not None: - if comm.strip() == name: - return True - else: - return False - - return False + status = _read_text("/proc/self/status") + if status is None: + return None + for line in status.splitlines(): + if line.startswith("NSpid:"): + pids = line.split()[1:] + # Single PID means container PID == host PID + return len(pids) == 1 + return None def detect_pid_host() -> bool: - ic = in_container() - if not ic: + """Check if os.getpid() matches the PID that NVML will report. + + NVML always reports host-namespace PIDs. If this process lives in a + different PID namespace the match will fail and process-scoped memory + tracking won't work. We detect this via the NSpid field in + /proc/self/status: a single entry means host PID == container PID. + """ + if not in_container(): return True - return has_pid_host() + result = _nvml_pid_matches() + if result is not None: + return result + + # Cannot determine — assume mismatch to be safe. + return False diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index 90f243151..802f8dd24 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -17,7 +17,7 @@ stage_args: worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.9 - enforce_eager: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" @@ -48,7 +48,7 @@ stage_args: worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 - enforce_eager: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav enable_prefix_caching: false diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml index a1b082b87..da996494e 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml @@ -86,7 +86,7 @@ stage_args: engine_output_type: audio # Final output: audio waveform gpu_memory_utilization: 0.1 distributed_executor_backend: "mp" - max_num_batched_tokens: 10000 + max_num_batched_tokens: 51200 # [TODO] if max_num_batch_tokens < max_batch_size * 800, there will be precision problem. hf_config_name: thinker_config engine_input_source: [1] final_output: true diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 701923309..11aea4f04 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -84,7 +84,7 @@ def _validate_stage_inputs(stage_list, engine_input_source): def thinker2talker_async_chunk( - connector: Any, + transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, ) -> list[dict[str, Any]]: @@ -96,7 +96,7 @@ def thinker2talker_async_chunk( """ request_id = request.external_req_id - chunk_id = connector.put_requests[request_id] + chunk_id = transfer_manager.put_req_chunk[request_id] if chunk_id == 0: all_token_ids = request.all_token_ids # prefill + decode prompt_token_ids = request.prompt_token_ids @@ -114,12 +114,12 @@ def thinker2talker_async_chunk( "tts_pad_embed": pooling_output.get("tts_pad_embed").detach().cpu(), "finished": torch.tensor(request.is_finished(), dtype=torch.bool), } - if connector.request_payload.get(request_id) is None: + if transfer_manager.request_payload.get(request_id) is None: if not request.is_finished(): - connector.request_payload[request_id] = talker_additional_info + transfer_manager.request_payload[request_id] = talker_additional_info return None else: - save_payload = connector.request_payload.pop(request_id) + save_payload = transfer_manager.request_payload.pop(request_id) talker_additional_info["thinker_embeddings"] = torch.cat( (save_payload.get("thinker_embeddings"), talker_additional_info.get("thinker_embeddings")), dim=0 ) @@ -206,7 +206,7 @@ def thinker2talker( def talker2code2wav_async_chunk( - connector: Any, + transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, ): @@ -241,8 +241,8 @@ def talker2code2wav_async_chunk( request_id = request.external_req_id chunk_size = left_context_size = 25 - connector.code_prompt_token_ids[request_id].append(codec_codes) - length = len(connector.code_prompt_token_ids[request_id]) + transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + length = len(transfer_manager.code_prompt_token_ids[request_id]) chunk_length = length % chunk_size if chunk_length != 0 and not request.is_finished(): return None @@ -252,7 +252,10 @@ def talker2code2wav_async_chunk( info = { "code_predictor_codes": ( - torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:]).transpose(0, 1).reshape(-1).tolist() + torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]) + .transpose(0, 1) + .reshape(-1) + .tolist() ), "finished": torch.tensor(request.is_finished(), dtype=torch.bool), } diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 90c718d40..e2f069d8f 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,5 +1,6 @@ import sys +from aenum import extend_enum from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, @@ -8,6 +9,7 @@ from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest from vllm.v1.request import Request as _OriginalRequest +from vllm.v1.request import RequestStatus import vllm_omni.logger # noqa: F401 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest @@ -15,6 +17,12 @@ from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding from vllm_omni.request import OmniRequest +# Extend RequestStatus enum with omni-specific statuses +if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"): + # The value - 1 is intentionally chosen to ensure it is treated + # as a non-finished state and remains compatible with existing comparisons. + extend_enum(RequestStatus, "WAITING_FOR_CHUNK", -1) + for module_name, module in sys.modules.items(): # only do patch on module of vllm, pass others if "vllm" not in module_name: