From c9392dff01982eeec9a8bd24fcc75f6a1a7fd9b2 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 28 Jan 2026 22:12:52 +0800 Subject: [PATCH] [Hardware] Support platforms and plugin system (#774) Signed-off-by: gcanlin --- .buildkite/pipeline.yml | 20 +- .buildkite/test-amd.yaml | 2 +- docs/configuration/stage_configs.md | 6 +- .../stage_configs/qwen2_5_omni.yaml | 6 +- .../image_to_image/image_edit.py | 22 +- .../image_to_image/image_to_image.md | 4 + .../image_to_video/README.md | 4 + .../image_to_video/image_to_video.py | 23 +- .../text_to_audio/text_to_audio.py | 5 +- .../offline_inference/text_to_image/README.md | 4 + .../text_to_image/gradio_demo.py | 9 +- .../text_to_image/text_to_image.py | 23 +- .../text_to_video/text_to_video.md | 4 + .../text_to_video/text_to_video.py | 23 +- .../qwen3_omni/qwen3_omni_moe_thinking.yaml | 2 +- .../diffusion/attention/test_attention_sp.py | 21 +- tests/diffusion/distributed/test_comm.py | 28 +- ...ion_worker.py => test_diffusion_worker.py} | 62 ++- .../stage_configs/npu/qwen2_5_omni_ci.yaml | 6 +- .../stage_configs/qwen2_5_omni_ci.yaml | 6 +- .../stage_configs/qwen3_omni_ci.yaml | 6 +- .../stage_configs/rocm/qwen2_5_omni_ci.yaml | 6 +- .../stage_configs/rocm/qwen3_omni_ci.yaml | 6 +- .../test_diffusion_cpu_offload.py | 6 +- .../offline_inference/test_diffusion_lora.py | 5 - .../offline_inference/test_qwen2_5_omni.py | 6 +- .../e2e/offline_inference/test_qwen3_omni.py | 4 +- .../test_sequence_parallel.py | 31 +- tests/e2e/offline_inference/test_t2i_model.py | 9 +- .../test_zimage_tensor_parallel.py | 4 +- .../stage_configs/qwen3_omni_ci.yaml | 10 +- .../stage_configs/qwen3_omni_thinker_ci.yaml | 2 +- .../stage_configs/rocm/qwen3_omni_ci.yaml | 6 +- .../test_images_generations_lora.py | 4 - tests/e2e/online_serving/test_qwen3_omni.py | 12 +- tests/e2e/stage_configs/qwen3_omni_ci.yaml | 6 +- tests/entrypoints/test_stage_utils.py | 55 +- .../diffusion/attention/backends/abstract.py | 56 +- .../attention/backends/ascend_attn.py | 92 ---- .../attention/backends/flash_attn.py | 54 +- .../diffusion/attention/backends/registry.py | 152 ++++++ .../diffusion/attention/backends/sage_attn.py | 2 +- .../diffusion/attention/backends/sdpa.py | 4 + .../diffusion/attention/backends/utils/fa.py | 4 +- vllm_omni/diffusion/attention/selector.py | 127 ++--- vllm_omni/diffusion/distributed/comm.py | 10 +- .../distributed/group_coordinator.py | 14 +- .../diffusion/distributed/parallel_state.py | 18 +- vllm_omni/diffusion/distributed/utils.py | 7 +- vllm_omni/diffusion/envs.py | 163 +----- .../diffusion/executor/multiproc_executor.py | 7 +- .../diffusion/hooks/sequence_parallel.py | 7 +- vllm_omni/diffusion/layers/adalayernorm.py | 8 + vllm_omni/diffusion/layers/custom_op.py | 14 +- vllm_omni/diffusion/layers/rope.py | 8 + .../longcat_image_transformer.py | 4 +- vllm_omni/diffusion/worker/__init__.py | 10 +- ...el_runner.py => diffusion_model_runner.py} | 2 +- ...iffusion_worker.py => diffusion_worker.py} | 30 +- vllm_omni/diffusion/worker/npu/__init__.py | 7 - vllm_omni/diffusion/worker/npu/npu_worker.py | 255 --------- vllm_omni/engine/arg_utils.py | 13 +- vllm_omni/entrypoints/omni_stage.py | 504 ++++++++++-------- vllm_omni/entrypoints/stage_utils.py | 6 +- vllm_omni/entrypoints/utils.py | 16 +- .../models/qwen2_5_omni/qwen2_5_omni.py | 4 +- .../qwen2_5_omni/qwen2_5_omni_token2wav.py | 8 +- .../models/qwen3_omni/qwen3_omni.py | 6 +- .../model_executor/stage_configs/bagel.yaml | 2 +- .../stage_configs/bagel_multiconnector.yaml | 2 +- .../stage_configs/qwen2_5_omni.yaml | 6 +- .../qwen2_5_omni_multiconnector.yaml | 6 +- .../stage_configs/qwen3_omni_moe.yaml | 6 +- .../qwen3_omni_moe_async_chunk.yaml | 6 +- .../qwen3_omni_moe_multiconnector.yaml | 6 +- .../stage_configs/qwen3_tts.yaml | 2 +- vllm_omni/platforms/__init__.py | 191 +++++++ vllm_omni/platforms/cuda/__init__.py | 6 + vllm_omni/platforms/cuda/platform.py | 101 ++++ vllm_omni/platforms/interface.py | 90 ++++ vllm_omni/platforms/npu/__init__.py | 6 + vllm_omni/platforms/npu/platform.py | 84 +++ .../npu/stage_configs}/qwen2_5_omni.yaml | 12 +- .../npu/stage_configs}/qwen3_omni_moe.yaml | 6 +- .../npu/stage_configs}/qwen3_tts.yaml | 2 +- .../npu => platforms/npu/worker}/__init__.py | 0 .../npu/worker}/npu_ar_model_runner.py | 2 +- .../npu/worker}/npu_ar_worker.py | 5 +- .../worker}/npu_generation_model_runner.py | 4 +- .../npu/worker}/npu_generation_worker.py | 5 +- .../npu/worker}/npu_model_runner.py | 0 vllm_omni/platforms/rocm/__init__.py | 6 + vllm_omni/platforms/rocm/platform.py | 97 ++++ .../rocm/stage_configs}/qwen2_5_omni.yaml | 6 +- .../rocm/stage_configs}/qwen3_omni_moe.yaml | 6 +- vllm_omni/platforms/xpu/__init__.py | 6 + vllm_omni/platforms/xpu/platform.py | 76 +++ vllm_omni/plugins/__init__.py | 81 +++ vllm_omni/utils/__init__.py | 13 - vllm_omni/utils/platform_utils.py | 58 -- vllm_omni/worker/gpu_ar_worker.py | 3 +- vllm_omni/worker/gpu_generation_worker.py | 3 +- vllm_omni/worker/mixins.py | 13 + 103 files changed, 1770 insertions(+), 1218 deletions(-) rename tests/diffusion/{test_gpu_diffusion_worker.py => test_diffusion_worker.py} (81%) delete mode 100644 vllm_omni/diffusion/attention/backends/ascend_attn.py create mode 100644 vllm_omni/diffusion/attention/backends/registry.py rename vllm_omni/diffusion/worker/{gpu_diffusion_model_runner.py => diffusion_model_runner.py} (99%) rename vllm_omni/diffusion/worker/{gpu_diffusion_worker.py => diffusion_worker.py} (94%) delete mode 100644 vllm_omni/diffusion/worker/npu/__init__.py delete mode 100644 vllm_omni/diffusion/worker/npu/npu_worker.py create mode 100644 vllm_omni/platforms/__init__.py create mode 100644 vllm_omni/platforms/cuda/__init__.py create mode 100644 vllm_omni/platforms/cuda/platform.py create mode 100644 vllm_omni/platforms/interface.py create mode 100644 vllm_omni/platforms/npu/__init__.py create mode 100644 vllm_omni/platforms/npu/platform.py rename vllm_omni/{model_executor/stage_configs/npu => platforms/npu/stage_configs}/qwen2_5_omni.yaml (89%) rename vllm_omni/{model_executor/stage_configs/npu => platforms/npu/stage_configs}/qwen3_omni_moe.yaml (93%) rename vllm_omni/{model_executor/stage_configs/npu => platforms/npu/stage_configs}/qwen3_tts.yaml (89%) rename vllm_omni/{worker/npu => platforms/npu/worker}/__init__.py (100%) rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_ar_model_runner.py (99%) rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_ar_worker.py (74%) rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_generation_model_runner.py (99%) rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_generation_worker.py (72%) rename vllm_omni/{worker/npu => platforms/npu/worker}/npu_model_runner.py (100%) create mode 100644 vllm_omni/platforms/rocm/__init__.py create mode 100644 vllm_omni/platforms/rocm/platform.py rename vllm_omni/{model_executor/stage_configs/rocm => platforms/rocm/stage_configs}/qwen2_5_omni.yaml (94%) rename vllm_omni/{model_executor/stage_configs/rocm => platforms/rocm/stage_configs}/qwen3_omni_moe.yaml (93%) create mode 100644 vllm_omni/platforms/xpu/__init__.py create mode 100644 vllm_omni/platforms/xpu/platform.py create mode 100644 vllm_omni/plugins/__init__.py delete mode 100644 vllm_omni/utils/platform_utils.py create mode 100644 vllm_omni/worker/mixins.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d1e71aba2..78a400cd8 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -18,11 +18,23 @@ steps: # queue: "ascend" - label: "Simple Unit Test" - depends_on: ~ + depends_on: image-build commands: - - ".buildkite/scripts/simple_test.sh" + - pytest -v -s tests/entrypoints/ + - pytest -v -s tests/diffusion/cache/ + - pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py + - pytest -v -s tests/worker/ agents: - queue: "cpu_queue_premerge" + queue: "gpu_1_queue" + 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 Model Test" timeout_in_minutes: 20 @@ -149,7 +161,7 @@ steps: timeout_in_minutes: 20 depends_on: image-build commands: - - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py + - pytest -s -v tests/diffusion/test_diffusion_worker.py agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 86d65f15b..08a3f7a1c 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -54,7 +54,7 @@ steps: commands: - export MIOPEN_DEBUG_CONV_DIRECT=0 - export MIOPEN_DEBUG_CONV_GEMM=0 - - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py + - pytest -s -v tests/diffusion/test_diffusion_worker.py - label: "Omni Model Test Qwen2-5-Omni" timeout_in_minutes: 15 diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 18972eb79..1f1a18902 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -40,7 +40,7 @@ stage_args: engine_args: # Engine arguments for a certain engine model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used + worker_type: ar # The specific worker used scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip enforce_eager: true # Now we only support eager mode @@ -66,7 +66,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true @@ -92,7 +92,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/docs/configuration/stage_configs/qwen2_5_omni.yaml b/docs/configuration/stage_configs/qwen2_5_omni.yaml index e71f348de..e20e79c3e 100644 --- a/docs/configuration/stage_configs/qwen2_5_omni.yaml +++ b/docs/configuration/stage_configs/qwen2_5_omni.yaml @@ -8,7 +8,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true # Now we only support eager mode @@ -34,7 +34,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true @@ -60,7 +60,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index f2065e6b2..896c611b4 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -81,7 +81,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import detect_device_type, is_npu +from vllm_omni.platforms import current_omni_platform def parse_args() -> argparse.Namespace: @@ -280,6 +280,16 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Disable torch.compile and force eager execution.", ) + parser.add_argument( + "--vae_use_slicing", + action="store_true", + help="Enable VAE slicing for memory optimization.", + ) + parser.add_argument( + "--vae_use_tiling", + action="store_true", + help="Enable VAE tiling for memory optimization.", + ) parser.add_argument( "--enable-cpu-offload", action="store_true", @@ -306,12 +316,8 @@ def main(): else: input_image = input_images - device = detect_device_type() - generator = torch.Generator(device=device).manual_seed(args.seed) + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - # Enable VAE memory optimizations on NPU - vae_use_slicing = is_npu() - vae_use_tiling = is_npu() parallel_config = DiffusionParallelConfig( ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, @@ -344,8 +350,8 @@ def main(): # Initialize Omni with appropriate pipeline omni = Omni( model=args.model, - vae_use_slicing=vae_use_slicing, - vae_use_tiling=vae_use_tiling, + 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, diff --git a/examples/offline_inference/image_to_image/image_to_image.md b/examples/offline_inference/image_to_image/image_to_image.md index d624b4684..b9e25bf98 100644 --- a/examples/offline_inference/image_to_image/image_to_image.md +++ b/examples/offline_inference/image_to_image/image_to_image.md @@ -47,4 +47,8 @@ Key arguments: - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models. - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower). - `--output`: path to save the generated PNG. +- `--vae_use_slicing`: enable VAE slicing for memory optimization. +- `--vae_use_tiling`: enable VAE tiling for memory optimization. - `--enable-cpu-offload`: enable CPU offloading for diffusion models. + +> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage. diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index afebb1963..52bb389e7 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -54,4 +54,8 @@ Key arguments: - `--num_inference_steps`: Number of denoising steps (default 50). - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video). - `--output`: Path to save the generated video. +- `--vae_use_slicing`: Enable VAE slicing for memory optimization. +- `--vae_use_tiling`: Enable VAE tiling for memory optimization. - `--enable-cpu-offload`: enable CPU offloading for diffusion models. + +> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage. 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 06c13fbee..4fca7825c 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -29,7 +29,7 @@ import torch from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import detect_device_type, is_npu +from vllm_omni.platforms import current_omni_platform def parse_args() -> argparse.Namespace: @@ -59,6 +59,16 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).") parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.") + parser.add_argument( + "--vae_use_slicing", + action="store_true", + help="Enable VAE slicing for memory optimization.", + ) + parser.add_argument( + "--vae_use_tiling", + action="store_true", + help="Enable VAE tiling for memory optimization.", + ) parser.add_argument( "--enable-cpu-offload", action="store_true", @@ -80,8 +90,7 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t def main(): args = parse_args() - device = detect_device_type() - generator = torch.Generator(device=device).manual_seed(args.seed) + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) # Load input image image = PIL.Image.open(args.image).convert("RGB") @@ -98,17 +107,13 @@ def main(): # Resize image to target dimensions image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) - # Enable VAE memory optimizations on NPU - vae_use_slicing = is_npu() - vae_use_tiling = is_npu() - # Check if profiling is requested via environment variable profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) omni = Omni( model=args.model, - vae_use_slicing=vae_use_slicing, - vae_use_tiling=vae_use_tiling, + vae_use_slicing=args.vae_use_slicing, + vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, flow_shift=args.flow_shift, enable_cpu_offload=args.enable_cpu_offload, diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 1124dbdb4..0a9efcca5 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -22,7 +22,7 @@ import torch from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.utils.platform_utils import detect_device_type +from vllm_omni.platforms import current_omni_platform def parse_args() -> argparse.Namespace: @@ -118,8 +118,7 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410 def main(): args = parse_args() - device = detect_device_type() - generator = torch.Generator(device=device).manual_seed(args.seed) + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) print(f"\n{'=' * 60}") print("Stable Audio Open - Text-to-Audio Generation") diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index f4de19db8..ab28b115c 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -96,8 +96,12 @@ Key arguments: - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower). - `--height/--width`: output resolution (defaults 1024x1024). - `--output`: path to save the generated PNG. +- `--vae_use_slicing`: enable VAE slicing for memory optimization. +- `--vae_use_tiling`: enable VAE tiling for memory optimization. - `--enable-cpu-offload`: enable CPU offloading for diffusion models. +> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage. + > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes. ## Web UI Demo diff --git a/examples/offline_inference/text_to_image/gradio_demo.py b/examples/offline_inference/text_to_image/gradio_demo.py index 30446f50e..614e40903 100644 --- a/examples/offline_inference/text_to_image/gradio_demo.py +++ b/examples/offline_inference/text_to_image/gradio_demo.py @@ -7,7 +7,7 @@ import torch from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import detect_device_type, is_npu +from vllm_omni.platforms import current_omni_platform ASPECT_RATIOS: dict[str, tuple[int, int]] = { "1:1": (1328, 1328), @@ -62,8 +62,8 @@ def parse_args() -> argparse.Namespace: @lru_cache(maxsize=1) def get_omni(model_name: str) -> Omni: # Enable VAE memory optimizations on NPU - vae_use_slicing = is_npu() - vae_use_tiling = is_npu() + vae_use_slicing = current_omni_platform.is_npu() + vae_use_tiling = current_omni_platform.is_npu() return Omni( model=model_name, vae_use_slicing=vae_use_slicing, @@ -72,7 +72,6 @@ def get_omni(model_name: str) -> Omni: def build_demo(args: argparse.Namespace) -> gr.Blocks: - device = detect_device_type() omni = get_omni(args.model) def run_inference( @@ -99,7 +98,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks: raise gr.Error("Inference steps must be a positive integer.") if num_images not in {1, 2, 3, 4}: raise gr.Error("Number of images must be 1, 2, 3, or 4.") - generator = torch.Generator(device=device).manual_seed(seed) + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed) outputs = omni.generate( prompt.strip(), OmniDiffusionSamplingParams( 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 efc5010e0..3e0628a15 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -12,7 +12,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import detect_device_type, is_npu +from vllm_omni.platforms import current_omni_platform def parse_args() -> argparse.Namespace: @@ -113,17 +113,22 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", ) + parser.add_argument( + "--vae_use_slicing", + action="store_true", + help="Enable VAE slicing for memory optimization.", + ) + parser.add_argument( + "--vae_use_tiling", + action="store_true", + help="Enable VAE tiling for memory optimization.", + ) return parser.parse_args() def main(): args = parse_args() - device = detect_device_type() - generator = torch.Generator(device=device).manual_seed(args.seed) - - # Enable VAE memory optimizations on NPU - vae_use_slicing = is_npu() - vae_use_tiling = is_npu() + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) # Configure cache based on backend type cache_config = None @@ -167,8 +172,8 @@ def main(): omni = Omni( model=args.model, - vae_use_slicing=vae_use_slicing, - vae_use_tiling=vae_use_tiling, + vae_use_slicing=args.vae_use_slicing, + vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, cache_config=cache_config, enable_cache_dit_summary=args.enable_cache_dit_summary, diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 7204ec58a..cd5bd2be4 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -29,4 +29,8 @@ Key arguments: - `--boundary_ratio`: Boundary split ratio for low/high DiT. - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video). - `--output`: path to save the generated video. +- `--vae_use_slicing`: enable VAE slicing for memory optimization. +- `--vae_use_tiling`: enable VAE tiling for memory optimization. - `--enable-cpu-offload`: enable CPU offloading for diffusion models. + +> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage. 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 49522278f..464ff9cca 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -13,7 +13,7 @@ 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.utils.platform_utils import detect_device_type, is_npu +from vllm_omni.platforms import current_omni_platform def parse_args() -> argparse.Namespace: @@ -54,6 +54,16 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).") parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.") + parser.add_argument( + "--vae_use_slicing", + action="store_true", + help="Enable VAE slicing for memory optimization.", + ) + parser.add_argument( + "--vae_use_tiling", + action="store_true", + help="Enable VAE tiling for memory optimization.", + ) parser.add_argument( "--enforce_eager", action="store_true", @@ -81,12 +91,7 @@ def parse_args() -> argparse.Namespace: def main(): args = parse_args() - device = detect_device_type() - generator = torch.Generator(device=device).manual_seed(args.seed) - - # Enable VAE memory optimizations on NPU - vae_use_slicing = is_npu() - vae_use_tiling = is_npu() + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment). cache_config = None @@ -118,8 +123,8 @@ def main(): omni = Omni( model=args.model, - vae_use_slicing=vae_use_slicing, - vae_use_tiling=vae_use_tiling, + vae_use_slicing=args.vae_use_slicing, + vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, flow_shift=args.flow_shift, cache_backend=args.cache_backend, diff --git a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml index a69bd2fe1..34b63e26e 100644 --- a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml +++ b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml @@ -13,7 +13,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.9 enforce_eager: true diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index c0061daf9..447d1ea93 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/tests/diffusion/attention/test_attention_sp.py @@ -6,7 +6,6 @@ import tempfile import pytest import torch -from vllm.platforms import current_platform from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import ( @@ -19,15 +18,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( initialize_model_parallel, ) from vllm_omni.diffusion.forward_context import set_forward_context -from vllm_omni.utils.platform_utils import detect_device_type - -device_type = detect_device_type() -if device_type == "cuda": - torch_device = torch.cuda -elif device_type == "npu": - torch_device = torch.npu -else: - raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.") +from vllm_omni.platforms import current_omni_platform def update_environment_variables(envs_dict: dict[str, str]): @@ -171,7 +162,7 @@ def test_sequence_parallel( sequence_parallel_size = ulysses_degree * ring_degree # Skip if not enough GPUs available - available_gpus = torch_device.device_count() + available_gpus = current_omni_platform.get_device_count() if available_gpus < sequence_parallel_size: pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available") @@ -338,13 +329,13 @@ def ulysses_attention_on_test_model( """Run Ulysses attention test on a test model and save results for comparison.""" # Use fixed seed for reproducibility across baseline and SP runs RANDOM_SEED = 42 - current_platform.seed_everything(RANDOM_SEED) + current_omni_platform.seed_everything(RANDOM_SEED) mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})" print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}") - device = torch.device(f"{device_type}:{local_rank}") - torch_device.set_device(device) + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -423,7 +414,7 @@ def ulysses_attention_on_test_model( # Generate and save full input data with fixed seed # Reinitialize RNG to ensure reproducibility torch.manual_seed(42) - torch_device.manual_seed_all(42) + current_omni_platform.seed_everything(42) full_hidden_states = torch.randn( (batch_size, seq_len, hidden_size), dtype=dtype, diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 0e76f27b7..3e408c10a 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/tests/diffusion/distributed/test_comm.py @@ -14,15 +14,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) -from vllm_omni.utils.platform_utils import detect_device_type - -device_type = detect_device_type() -if device_type == "cuda": - torch_device = torch.cuda -elif device_type == "npu": - torch_device = torch.npu -else: - raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.") +from vllm_omni.platforms import current_omni_platform def update_environment_variables(envs_dict: dict[str, str]): @@ -49,7 +41,7 @@ def test_4d_identity( ): """Test that two consecutive all-to-all operations return the original input.""" # Skip if not enough GPUs available - available_gpus = torch_device.device_count() + available_gpus = current_omni_platform.get_device_count() if available_gpus < world_size: pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") @@ -85,8 +77,8 @@ def _test_4d_identity_worker( ): """Worker function for test_4d_identity.""" # Set device - device = torch.device(f"{device_type}:{local_rank}") - torch_device.set_device(device) + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) # Set environment variables for distributed training update_environment_variables( @@ -183,7 +175,7 @@ def test_5d_identity( ): """Test that two consecutive all-to-all operations return the original input.""" # Skip if not enough GPUs available - available_gpus = torch_device.device_count() + available_gpus = current_omni_platform.get_device_count() if available_gpus < world_size: pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") @@ -219,8 +211,8 @@ def _test_5d_identity_worker( ): """Worker function for test_5d_identity.""" # Set device - device = torch.device(f"{device_type}:{local_rank}") - torch_device.set_device(device) + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) # Set environment variables for distributed training update_environment_variables( @@ -316,7 +308,7 @@ def test_ring_p2p( ): """Test Ring P2P communication (send_recv).""" # Skip if not enough GPUs available - available_gpus = torch_device.device_count() + available_gpus = current_omni_platform.get_device_count() if available_gpus < world_size: pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") @@ -339,8 +331,8 @@ def _test_ring_p2p_worker( import sys # Set device - device = torch.device(f"{device_type}:{local_rank}") - torch_device.set_device(device) + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) # Set env vars # Use a different port to avoid conflict with other tests if run in parallel diff --git a/tests/diffusion/test_gpu_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py similarity index 81% rename from tests/diffusion/test_gpu_diffusion_worker.py rename to tests/diffusion/test_diffusion_worker.py index 7a43710c8..220f210a3 100644 --- a/tests/diffusion/test_gpu_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Unit tests for GPUDiffusionWorker class. +Unit tests for DiffusionWorker class. -This module tests the GPUDiffusionWorker implementation: +This module tests the DiffusionWorker implementation: - load_weights: Loading model weights - sleep: Putting worker into sleep mode (levels 1 and 2) - wake_up: Waking worker from sleep mode @@ -15,7 +15,7 @@ from unittest.mock import Mock, patch import pytest import torch -from vllm_omni.diffusion.worker.gpu_diffusion_worker import GPUDiffusionWorker +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker @pytest.fixture @@ -33,18 +33,19 @@ def mock_od_config(): @pytest.fixture def mock_gpu_worker(mock_od_config): - """Create a GPUDiffusionWorker with mocked initialization.""" - with patch.object(GPUDiffusionWorker, "init_device"): - worker = GPUDiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config) + """Create a DiffusionWorker with mocked initialization.""" + with patch.object(DiffusionWorker, "init_device"): + worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config) # Mock the model_runner with pipeline worker.model_runner = Mock() worker.model_runner.pipeline = Mock() + worker.device = torch.device("cuda", 0) worker._sleep_saved_buffers = {} return worker -class TestGPUDiffusionWorkerLoadWeights: - """Test GPUDiffusionWorker.load_weights method.""" +class TestDiffusionWorkerLoadWeights: + """Test DiffusionWorker.load_weights method.""" def test_load_weights_calls_pipeline(self, mock_gpu_worker): """Test that load_weights delegates to model_runner.load_weights.""" @@ -75,20 +76,21 @@ class TestGPUDiffusionWorkerLoadWeights: assert result == set() -class TestGPUDiffusionWorkerSleep: - """Test GPUDiffusionWorker.sleep method.""" +class TestDiffusionWorkerSleep: + """Test DiffusionWorker.sleep method.""" - @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info") + @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_mem_info, mock_gpu_worker): + def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" # Setup memory info mocks - # Before sleep: 1GB free, 8GB total - # After sleep: 3GB free, 8GB total (freed 2GB) - mock_mem_info.side_effect = [ - (1 * 1024**3, 8 * 1024**3), # Before sleep - (3 * 1024**3, 8 * 1024**3), # After sleep + # 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 ] + mock_platform.get_device_total_memory.return_value = 8 * 1024**3 # Setup allocator mock mock_allocator = Mock() @@ -104,15 +106,16 @@ class TestGPUDiffusionWorkerSleep: # 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.gpu_diffusion_worker.torch.cuda.mem_get_info") + @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_mem_info, mock_gpu_worker): + def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" # Setup memory info mocks - mock_mem_info.side_effect = [ - (1 * 1024**3, 8 * 1024**3), # Before sleep - (5 * 1024**3, 8 * 1024**3), # After sleep (freed 4GB) + mock_platform.get_free_memory.side_effect = [ + 1 * 1024**3, # Before sleep + 5 * 1024**3, # After sleep (freed 4GB) ] + mock_platform.get_device_total_memory.return_value = 8 * 1024**3 # Setup allocator mock mock_allocator = Mock() @@ -141,15 +144,16 @@ class TestGPUDiffusionWorkerSleep: assert "buffer1" in mock_gpu_worker._sleep_saved_buffers assert "buffer2" in mock_gpu_worker._sleep_saved_buffers - @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info") + @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_mem_info, mock_gpu_worker): + def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker): """Test that sleep validates memory was actually freed.""" # Simulate memory increase (should trigger assertion error) - mock_mem_info.side_effect = [ - (3 * 1024**3, 8 * 1024**3), # Before sleep: 3GB free - (1 * 1024**3, 8 * 1024**3), # After sleep: 1GB free (negative freed!) + mock_platform.get_free_memory.side_effect = [ + 3 * 1024**3, # Before sleep: 3GB free + 1 * 1024**3, # After sleep: 1GB free (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) @@ -160,8 +164,8 @@ class TestGPUDiffusionWorkerSleep: mock_gpu_worker.sleep(level=1) -class TestGPUDiffusionWorkerWakeUp: - """Test GPUDiffusionWorker.wake_up method.""" +class TestDiffusionWorkerWakeUp: + """Test DiffusionWorker.wake_up method.""" @patch("vllm.device_allocator.cumem.CuMemAllocator") def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker): diff --git a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml index d2aae1c0d..d6b04258a 100644 --- a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml @@ -10,7 +10,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler max_model_len: 896 max_num_batched_tokens: 896 @@ -40,7 +40,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler max_model_len: 896 max_num_batched_tokens: 896 @@ -70,7 +70,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml index eb9109adc..e093ec51b 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml @@ -11,7 +11,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler max_model_len: 896 max_num_batched_tokens: 896 @@ -41,7 +41,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler max_model_len: 896 max_num_batched_tokens: 896 @@ -71,7 +71,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml index 0780420fd..7ac435ea5 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml @@ -12,7 +12,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -42,7 +42,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: true @@ -73,7 +73,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml index 96e9d7fa7..474df5e79 100644 --- a/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml @@ -11,7 +11,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler max_model_len: 896 max_num_batched_tokens: 896 @@ -41,7 +41,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler max_model_len: 896 max_num_batched_tokens: 896 @@ -71,7 +71,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml index 0780420fd..7ac435ea5 100644 --- a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml @@ -12,7 +12,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -42,7 +42,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: true @@ -73,7 +73,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index f6606dce6..35e106df8 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -7,7 +7,7 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory from tests.utils import GPUMemoryMonitor from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.utils.platform_utils import is_npu, is_rocm +from vllm_omni.platforms import current_omni_platform # ruff: noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[2] @@ -20,7 +20,7 @@ models = ["riverclouds/qwen_image_random"] def inference(model_name: str, offload: bool = True): - torch.cuda.empty_cache() + current_omni_platform.empty_cache() device_index = torch.cuda.current_device() monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() @@ -45,7 +45,7 @@ def inference(model_name: str, offload: bool = True): return peak -@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported") +@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported") @pytest.mark.parametrize("model_name", models) def test_cpu_offload_diffusion_model(model_name: str): try: diff --git a/tests/e2e/offline_inference/test_diffusion_lora.py b/tests/e2e/offline_inference/test_diffusion_lora.py index ef734da61..2465bc607 100644 --- a/tests/e2e/offline_inference/test_diffusion_lora.py +++ b/tests/e2e/offline_inference/test_diffusion_lora.py @@ -9,7 +9,6 @@ from safetensors.torch import save_file from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import is_npu # ruff: noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[2] @@ -25,10 +24,6 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" # model to reduce runtime and avoid extra downloads. models = ["Tongyi-MAI/Z-Image-Turbo"] -# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly. -if is_npu(): - pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.", allow_module_level=True) - @pytest.mark.parametrize("model_name", models) def test_diffusion_model(model_name: str, tmp_path: Path): diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni.py b/tests/e2e/offline_inference/test_qwen2_5_omni.py index 09b0306c1..3537ea127 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni.py @@ -13,7 +13,7 @@ from vllm.assets.video import VideoAsset from vllm.envs import VLLM_USE_MODELSCOPE from vllm.multimodal.image import convert_image_mode -from vllm_omni.utils import is_npu, is_rocm +from vllm_omni.platforms import current_omni_platform from .conftest import OmniRunner from .utils import create_new_process_for_each_test @@ -21,9 +21,9 @@ from .utils import create_new_process_for_each_test models = ["Qwen/Qwen2.5-Omni-3B"] # CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU -if is_npu(): +if current_omni_platform.is_npu(): stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml") -elif is_rocm(): +elif current_omni_platform.is_rocm(): # ROCm stage config optimized for MI325 GPU stage_config = str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml") else: diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index edd5958a0..470cdb73b 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -14,14 +14,14 @@ from pathlib import Path import pytest from vllm.assets.video import VideoAsset -from vllm_omni.utils import is_rocm +from vllm_omni.platforms import current_omni_platform from .conftest import OmniRunner models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] # CI stage config for 2xH100-80G GPUs or AMD GPU MI325 -if is_rocm(): +if current_omni_platform.is_rocm(): # ROCm stage config optimized for MI325 GPU stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")] else: diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index 367e2e59e..262c80243 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/tests/e2e/offline_inference/test_sequence_parallel.py @@ -30,8 +30,7 @@ if str(REPO_ROOT) not in sys.path: from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.diffusion.distributed.parallel_state import device_count -from vllm_omni.diffusion.envs import get_device_name +from vllm_omni.platforms import current_omni_platform models = ["riverclouds/qwen_image_random"] @@ -63,9 +62,9 @@ def _cleanup_distributed(): gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + current_omni_platform.synchronize() time.sleep(5) @@ -88,7 +87,7 @@ def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height width=width, num_inference_steps=4, guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(seed), + generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed), num_outputs_per_prompt=1, ), ) @@ -125,7 +124,7 @@ def _run_sp( width=width, num_inference_steps=4, guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(seed), + generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed), num_outputs_per_prompt=1, ), ) @@ -157,8 +156,8 @@ def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str): @pytest.mark.parametrize("attn_backend", ["sdpa"]) def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str): """Test SP inference only (ulysses=2).""" - if device_count() < 2: - pytest.skip(f"Test requires 2 GPUs but only {device_count()} available") + if current_omni_platform.get_device_count() < 2: + pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available") height = 256 width = 256 @@ -177,8 +176,8 @@ def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str @pytest.mark.parametrize("attn_backend", ["sdpa"]) def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str): """Test SP inference only (ring=2).""" - if device_count() < 2: - pytest.skip(f"Test requires 2 GPUs but only {device_count()} available") + if current_omni_platform.get_device_count() < 2: + pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available") height = 256 width = 256 @@ -212,8 +211,8 @@ def test_sequence_parallel( ) sp_size = ulysses_degree * ring_degree - if device_count() < sp_size: - pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available") + if current_omni_platform.get_device_count() < sp_size: + pytest.skip(f"Test requires {sp_size} GPUs but only {current_omni_platform.get_device_count()} available") height = 256 width = 256 @@ -243,8 +242,10 @@ def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_ba ulysses_degree = 4 ring_degree = 1 - if device_count() < ulysses_degree * ring_degree: - pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available") + if current_omni_platform.get_device_count() < ulysses_degree * ring_degree: + pytest.skip( + f"Test requires {ulysses_degree * ring_degree} GPUs but only {current_omni_platform.get_device_count()} available" + ) height = 272 width = 272 diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py index ba6cac34c..e7351f157 100644 --- a/tests/e2e/offline_inference/test_t2i_model.py +++ b/tests/e2e/offline_inference/test_t2i_model.py @@ -7,7 +7,7 @@ import torch from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import is_npu, is_rocm +from vllm_omni.platforms import current_omni_platform # ruff: noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[2] @@ -21,12 +21,11 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"] -# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly # Modelscope can't find riverclouds/qwen_image_random # TODO: When NPU support is ready, remove this branch. -if is_npu(): - models = ["Qwen/Qwen-Image"] -elif is_rocm(): +if current_omni_platform.is_npu(): + models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"] +elif current_omni_platform.is_rocm(): # TODO: When ROCm support is ready, remove this branch. # vLLM V0.11.0 has issues running riverclouds/qwen_image_random # on ROCm diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py index ea8f0acfe..0e3b97ec3 100644 --- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py +++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py @@ -23,7 +23,7 @@ from tests.utils import GPUMemoryMonitor from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.platform_utils import is_npu, is_rocm +from vllm_omni.platforms import current_omni_platform # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -129,7 +129,7 @@ def _run_zimage_generate( @pytest.mark.integration def test_zimage_tensor_parallel_tp2(tmp_path: Path): - if is_npu() or is_rocm(): + if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") if not torch.cuda.is_available() or torch.cuda.device_count() < 2: pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.") diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml index c5987ebba..4db11e11f 100644 --- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml @@ -6,13 +6,14 @@ # The following config has been verified on 2x H100-80G GPUs. stage_args: - stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM runtime: devices: "0,1" max_batch_size: 5 engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -20,6 +21,7 @@ stage_args: engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" enable_prefix_caching: false + max_num_batched_tokens: 32768 hf_config_name: thinker_config tensor_parallel_size: 2 load_format: dummy @@ -36,13 +38,14 @@ stage_args: repetition_penalty: 1.05 - stage_id: 1 + stage_type: llm # Use llm stage type to launch OmniLLM runtime: devices: "1" max_batch_size: 5 engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: false @@ -67,13 +70,14 @@ stage_args: stop_token_ids: [2150] - stage_id: 2 + stage_type: llm # Use llm stage type to launch OmniLLM runtime: devices: "0" max_batch_size: 1 engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml index 64c6b6ee3..89d03966d 100644 --- a/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml @@ -7,7 +7,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: true diff --git a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml index ef0bae0fa..0488a06c9 100644 --- a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml @@ -12,7 +12,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -42,7 +42,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: true @@ -73,7 +73,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py index 82044cd52..11165c6d2 100644 --- a/tests/e2e/online_serving/test_images_generations_lora.py +++ b/tests/e2e/online_serving/test_images_generations_lora.py @@ -27,8 +27,6 @@ from PIL import Image from safetensors.torch import save_file from vllm.utils.network_utils import get_open_port -from vllm_omni.utils.platform_utils import is_npu - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "Tongyi-MAI/Z-Image-Turbo" @@ -122,8 +120,6 @@ class OmniServer: @pytest.fixture(scope="module") def omni_server(): - if is_npu(): - pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.") with OmniServer(MODEL, ["--num-gpus", "1"]) as server: yield server diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 8564104a0..eb134ab5e 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -27,7 +27,7 @@ from tests.conftest import ( generate_synthetic_video, modify_stage_config, ) -from vllm_omni.utils import is_rocm +from vllm_omni.platforms import current_omni_platform models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -57,7 +57,7 @@ def get_chunk_config(): CHUNK_CONFIG_PATH = get_chunk_config() # CI stage config for 2xH100-80G GPUs or AMD GPU MI325 -if is_rocm(): +if current_omni_platform.is_rocm(): # ROCm stage config optimized for MI325 GPU stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")] else: @@ -144,7 +144,9 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues") +@pytest.mark.skipif( + current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues" +) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None: """ @@ -224,7 +226,9 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N assert similarity > 0.9, "The audio content is not same as the text" -@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues") +@pytest.mark.skipif( + current_omni_platform.is_rocm(), reason="Test skipped on AMD environment due to known output issues" +) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None: """ diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml index 60343dfdd..d7b07ddc9 100644 --- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml @@ -12,7 +12,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -43,7 +43,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: false @@ -71,7 +71,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index 67a2550b7..ac503639b 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -1,5 +1,6 @@ import os import sys +from unittest.mock import MagicMock import pytest @@ -42,6 +43,14 @@ def _make_dummy_torch(call_log): return _Torch +def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"): + """Create a mock platform for testing.""" + mock_platform = MagicMock() + mock_platform.device_type = device_type + mock_platform.device_control_env_var = env_var + return mock_platform + + @pytest.mark.usefixtures("clean_gpu_memory_between_tests") def test_set_stage_devices_respects_logical_ids(monkeypatch): # Preserve an existing logical mapping and ensure devices "0,1" map through it. @@ -49,9 +58,51 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch): call_log: list[int] = [] dummy_torch = _make_dummy_torch(call_log) monkeypatch.setitem(sys.modules, "torch", dummy_torch) - monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda") - monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES") + + # Mock the platform at the source module where it's defined + mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES") + monkeypatch.setattr( + "vllm_omni.platforms.current_omni_platform", + mock_platform, + ) set_stage_devices(stage_id=0, devices="0,1") assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7" + + +@pytest.mark.usefixtures("clean_gpu_memory_between_tests") +def test_set_stage_devices_npu_platform(monkeypatch): + """Test that set_stage_devices works correctly for NPU platform.""" + monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5") + call_log: list[int] = [] + + # Create NPU mock torch + class _Npu: + @staticmethod + def is_available(): + return True + + @staticmethod + def set_device(idx): + call_log.append(idx) + + @staticmethod + def device_count(): + return 2 + + class _NpuTorch: + npu = _Npu + + monkeypatch.setitem(sys.modules, "torch", _NpuTorch) + + # Mock NPU platform at the source module where it's defined + mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES") + monkeypatch.setattr( + "vllm_omni.platforms.current_omni_platform", + mock_platform, + ) + + set_stage_devices(stage_id=0, devices="0,1") + + assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5" diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 54906efc9..d0a62bcd9 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -7,12 +7,18 @@ from typing import Generic, TypeVar import torch +from vllm_omni.platforms import current_omni_platform + class AttentionBackend(ABC): """Abstract class for diffusion attention backends.""" accept_output_buffer: bool = False + @classmethod + def supports_attention_mask(cls) -> bool: + return False + @staticmethod @abstractmethod def get_name() -> str: @@ -77,12 +83,58 @@ class AttentionImpl(ABC, Generic[T]): ) -> None: raise NotImplementedError - @abstractmethod def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attn_metadata: T, + attn_metadata: T | None = None, + ) -> torch.Tensor: + """Dispatch to platform-specific forward implementation.""" + if current_omni_platform.is_rocm(): + return self.forward_hip(query, key, value, attn_metadata) + elif current_omni_platform.is_cuda(): + return self.forward_cuda(query, key, value, attn_metadata) + elif current_omni_platform.is_npu(): + return self.forward_npu(query, key, value, attn_metadata) + elif current_omni_platform.is_xpu(): + return self.forward_xpu(query, key, value, attn_metadata) + else: + raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}") + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T | None = None, ) -> torch.Tensor: raise NotImplementedError + + def forward_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T | None = None, + ) -> torch.Tensor: + raise NotImplementedError + + def forward_xpu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T | None = None, + ) -> torch.Tensor: + raise NotImplementedError + + def forward_hip( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T | None = None, + ) -> torch.Tensor: + # By default, HIP ops are compatible with CUDA ops. + return self.forward_cuda(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/backends/ascend_attn.py b/vllm_omni/diffusion/attention/backends/ascend_attn.py deleted file mode 100644 index 6563385c9..000000000 --- a/vllm_omni/diffusion/attention/backends/ascend_attn.py +++ /dev/null @@ -1,92 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from importlib.util import find_spec - -import torch -from vllm.logger import init_logger - -from vllm_omni.diffusion.attention.backends.abstract import ( - AttentionBackend, - AttentionImpl, - AttentionMetadata, -) - -logger = init_logger(__name__) - - -class AscendAttentionBackend(AttentionBackend): - accept_output_buffer: bool = True - - @staticmethod - def get_supported_head_sizes() -> list[int]: - return [x for x in range(1024)] # todo - - @staticmethod - def get_name() -> str: - return "ASCEND" - - @staticmethod - def get_impl_cls() -> type["AscendAttentionBackendImpl"]: - return AscendAttentionBackendImpl - - -class AscendAttentionBackendImpl(AttentionImpl): - def __init__( - self, - num_heads: int, - head_size: int, - softmax_scale: float, - causal: bool = False, - num_kv_heads: int | None = None, - prefix: str = "", - **extra_impl_args, - ) -> None: - self.causal = causal - self.softmax_scale = softmax_scale - - def forward_native( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: AttentionMetadata = None, - ) -> torch.Tensor: - query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) - attention_mask = attn_metadata.attn_mask if attn_metadata else None - - output = torch.nn.functional.scaled_dot_product_attention( - query, - key, - value, - attn_mask=attention_mask, - dropout_p=0.0, - is_causal=self.causal, - scale=self.softmax_scale, - ) - out = output.permute(0, 2, 1, 3) - return out - - def forward( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: AttentionMetadata = None, - ) -> torch.Tensor: - if find_spec("mindiesd"): - from mindiesd import attention_forward - - attention_mask = attn_metadata.attn_mask if attn_metadata else None - output = attention_forward( - query, - key, - value, - attn_mask=attention_mask, - opt_mode="manual", - op_type="fused_attn_score", - layout="BNSD", - ) - return output - else: - return self.forward_native(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index e9102b1d3..a9f2d925c 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -9,13 +9,6 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) -from vllm_omni.diffusion.attention.backends.utils.fa import ( - _pad_input, - _unpad_input, - _upad_input, - flash_attn_func, - flash_attn_varlen_func, -) logger = init_logger(__name__) @@ -23,6 +16,10 @@ logger = init_logger(__name__) class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True + @classmethod + def supports_attention_mask(cls) -> bool: + return True + @staticmethod def get_supported_head_sizes() -> list[int]: return [64, 96, 128, 192, 256] @@ -51,25 +48,22 @@ class FlashAttentionImpl(AttentionImpl): self.causal = causal self.softmax_scale = softmax_scale - def forward( + def forward_cuda( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: - """ - Flash attention implementation. + """CUDA/ROCm flash attention implementation.""" + from vllm_omni.diffusion.attention.backends.utils.fa import ( + _pad_input, + _unpad_input, + _upad_input, + flash_attn_func, + flash_attn_varlen_func, + ) - Args: - query: (batch_size, seq_len, num_heads, head_dim) - key: (batch_size, seq_len, num_heads, head_dim) - value: (batch_size, seq_len, num_heads, head_dim) - attn_metadata: AttentionMetadata. Attention mask is supported as attn_metadata.attn_mask - - Returns: - (batch_size, seq_len, num_heads, head_dim) - """ query_length = query.size(1) attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None # Contains at least one padding token in the sequence @@ -106,3 +100,25 @@ class FlashAttentionImpl(AttentionImpl): softmax_scale=self.softmax_scale, ) return out + + def forward_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata = None, + ) -> torch.Tensor: + """NPU attention implementation using mindiesd.""" + from mindiesd import attention_forward + + attention_mask = attn_metadata.attn_mask if attn_metadata else None + output = attention_forward( + query, + key, + value, + attn_mask=attention_mask, + opt_mode="manual", + op_type="fused_attn_score", + layout="BNSD", + ) + return output diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py new file mode 100644 index 000000000..a77223d16 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diffusion attention backend registry. + +This module provides an enum-based registry for diffusion attention backends, +similar to vLLM's AttentionBackendEnum. Each backend registers its class path, +and platforms can override or extend backends using register_backend(). +""" + +from collections.abc import Callable +from enum import Enum, EnumMeta +from typing import TYPE_CHECKING + +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname + +if TYPE_CHECKING: + from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend + +logger = init_logger(__name__) + + +class _DiffusionBackendEnumMeta(EnumMeta): + """Metaclass for DiffusionAttentionBackendEnum to provide better error messages.""" + + def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum": + """Get backend by name with helpful error messages.""" + try: + return super().__getitem__(name) # type: ignore[return-value] + except KeyError: + members = list(cls.__members__.keys()) + valid_backends = ", ".join(members) + raise ValueError( + f"Unknown diffusion attention backend: '{name}'. Valid options are: {valid_backends}" + ) from None + + +class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta): + """Enumeration of all supported diffusion attention backends. + + The enum value is the default class path, but this can be overridden + at runtime using register_backend(). + + To get the actual backend class (respecting overrides), use: + backend.get_class() + + Example: + # Get backend class + backend = DiffusionAttentionBackendEnum.FLASH_ATTN + backend_cls = backend.get_class() + + # Register custom backend + @register_diffusion_backend(DiffusionAttentionBackendEnum.CUSTOM) + class MyCustomBackend: + ... + """ + + # Common backends (available on most platforms) + FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend" + TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend" + SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend" + + def get_path(self, include_classname: bool = True) -> str: + """Get the class path for this backend (respects overrides). + + Returns: + The fully qualified class path string + + Raises: + ValueError: If backend has empty path and is not registered + """ + path = _DIFFUSION_ATTN_OVERRIDES.get(self, self.value) + if not path: + raise ValueError( + f"Backend {self.name} must be registered before use. " + f"Use register_diffusion_backend(DiffusionAttentionBackendEnum.{self.name}, " + f"'your.module.YourClass')" + ) + if not include_classname: + path = path.rsplit(".", 1)[0] + return path + + def get_class(self) -> "type[AttentionBackend]": + """Get the backend class (respects overrides). + + Returns: + The backend class + + Raises: + ImportError: If the backend class cannot be imported + ValueError: If backend has empty path and is not registered + """ + return resolve_obj_by_qualname(self.get_path()) + + def is_overridden(self) -> bool: + """Check if this backend has been overridden. + + Returns: + True if the backend has a registered override + """ + return self in _DIFFUSION_ATTN_OVERRIDES + + def clear_override(self) -> None: + """Clear any override for this backend, reverting to the default.""" + _DIFFUSION_ATTN_OVERRIDES.pop(self, None) + + +# Override registry +_DIFFUSION_ATTN_OVERRIDES: dict[DiffusionAttentionBackendEnum, str] = {} + + +def register_diffusion_backend( + backend: DiffusionAttentionBackendEnum, + class_path: str | None = None, +) -> Callable[[type], type]: + """Register or override a diffusion backend implementation. + + Args: + backend: The DiffusionAttentionBackendEnum member to register + class_path: Optional class path. If not provided and used as + decorator, will be auto-generated from the class. + + Returns: + Decorator function if class_path is None, otherwise a no-op + + Examples: + # Override an existing backend + @register_diffusion_backend(DiffusionAttentionBackendEnum.FLASH_ATTN) + class MyCustomFlashAttn: + ... + + # Override an existing backend (e.g., ASCEND_ATTN) + @register_diffusion_backend(DiffusionAttentionBackendEnum.ASCEND_ATTN) + class CustomAscendAttentionBackend: + ... + + # Direct registration + register_diffusion_backend( + DiffusionAttentionBackendEnum.CUSTOM, + "my.module.MyCustomBackend" + ) + """ + + def decorator(cls: type) -> type: + _DIFFUSION_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" + return cls + + if class_path is not None: + _DIFFUSION_ATTN_OVERRIDES[backend] = class_path + return lambda x: x + + return decorator diff --git a/vllm_omni/diffusion/attention/backends/sage_attn.py b/vllm_omni/diffusion/attention/backends/sage_attn.py index bfc06a839..af34a0633 100644 --- a/vllm_omni/diffusion/attention/backends/sage_attn.py +++ b/vllm_omni/diffusion/attention/backends/sage_attn.py @@ -54,7 +54,7 @@ class SageAttentionImpl(AttentionImpl): self.causal = causal self.softmax_scale = softmax_scale - def forward( + def forward_cuda( self, query: torch.Tensor, key: torch.Tensor, diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py index 89b3e84a7..16bc5fcf8 100644 --- a/vllm_omni/diffusion/attention/backends/sdpa.py +++ b/vllm_omni/diffusion/attention/backends/sdpa.py @@ -16,6 +16,10 @@ logger = init_logger(__name__) class SDPABackend(AttentionBackend): accept_output_buffer: bool = True + @classmethod + def supports_attention_mask(cls) -> bool: + return True + @staticmethod def get_supported_head_sizes() -> list[int]: return [x for x in range(1024)] # todo diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index 6a02435a7..ea12bc0f1 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -15,9 +15,9 @@ import torch import torch.nn.functional as F -from vllm_omni.utils.platform_utils import is_rocm +from vllm_omni.platforms import current_omni_platform -if is_rocm(): +if current_omni_platform.is_rocm(): from vllm._aiter_ops import is_aiter_found_and_supported # Choose to enable this by default on ROCm diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index ed5946098..920396e66 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/vllm_omni/diffusion/attention/selector.py @@ -1,52 +1,54 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Diffusion attention backend selector. + +This module provides the interface for selecting diffusion attention backends. +The actual backend selection logic is delegated to the platform layer +(vllm_omni.platforms), similar to how vLLM handles attention backend selection. + +Usage: + from vllm_omni.diffusion.attention.selector import get_attn_backend + + # Get the appropriate backend for current platform + backend_cls = get_attn_backend(head_size=64) + + # Or override via environment variable + # export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN +""" import importlib import os from functools import cache -import torch from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.abstract import ( AttentionBackend, ) -from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend -from vllm_omni.utils.platform_utils import detect_device_type, is_rocm logger = init_logger(__name__) -# environment variable value -> backend module and class -_BACKEND_CONFIG = { - "FLASH_ATTN": { - "module": "vllm_omni.diffusion.attention.backends.flash_attn", - "class": "FlashAttentionBackend", - }, - "TORCH_SDPA": { - "module": "vllm_omni.diffusion.attention.backends.sdpa", - "class": "SDPABackend", - }, - "SAGE_ATTN": { - "module": "vllm_omni.diffusion.attention.backends.sage_attn", - "class": "SageAttentionBackend", - }, - "ASCEND": {"module": "vllm_omni.diffusion.attention.backends.ascend_attn", "class": "AscendAttentionBackend"}, -} -_BACKENDS_SUPPORT_ATTENTION_MASK = ["SDPA", "ASCEND", "FLASH_ATTN"] +def _load_backend_cls(cls_path: str) -> type[AttentionBackend]: + """Load a backend class from its fully qualified path. + Args: + cls_path: Fully qualified class path (e.g., + "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend") -def load_backend(backend_name: str) -> type[AttentionBackend]: - config = _BACKEND_CONFIG[backend_name] - + Returns: + The loaded backend class + """ + module_path, class_name = cls_path.rsplit(".", 1) try: - module = importlib.import_module(config["module"]) - backend_class = getattr(module, config["class"]) + module = importlib.import_module(module_path) + backend_class = getattr(module, class_name) return backend_class except ImportError as e: - raise ImportError(f"Failed to import module {config['module']}: {e}") + raise ImportError(f"Failed to import module {module_path}: {e}") except AttributeError as e: - raise AttributeError(f"Class {config['class']} not found in module: {e}") + raise AttributeError(f"Class {class_name} not found in module: {e}") @cache @@ -54,65 +56,30 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]: """ Get attention backend for diffusion models. - The backend is selected based on the following priority: - 1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN) - 2. Default backend (SDPA) + The backend selection is delegated to the current platform + (vllm_omni.platforms.current_omni_platform), which selects the + appropriate backend based on: + 1. User override via DIFFUSION_ATTENTION_BACKEND environment variable + 2. Platform-specific defaults and capabilities + + This is similar to how vLLM's get_attn_backend_cls works, where the + platform layer decides which backend to use based on hardware capabilities. Args: - head_size: Head size (currently not used for selection, but kept for API compatibility) + head_size: Head size for attention computation (may affect backend selection) Returns: The selected attention backend class """ - # Check environment variable + from vllm_omni.platforms import current_omni_platform - backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND", None) + # Check environment variable for user override + selected_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND") - if detect_device_type() == "cuda" and not is_rocm(): - compute_capability = torch.cuda.get_device_capability() - major, minor = compute_capability - if 80 <= major * 10 + minor < 100: - if backend_name is None: - backend_name = "FLASH_ATTN" - else: - if backend_name == "FLASH_ATTN": - logger.warning( - """Flash Attention requires GPU with compute capability >= 8.0 or < 10.0. " - "Falling back to TORCH_SDPA backend.""" - ) - backend_name = "TORCH_SDPA" - elif detect_device_type() == "cuda" and is_rocm(): - from vllm._aiter_ops import is_aiter_found_and_supported + # Delegate to platform for backend selection + backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls( + selected_backend=selected_backend, + head_size=head_size, + ) - compute_capability = torch.cuda.get_device_capability() - major, minor = compute_capability - - # Choose to enable this by default on ROCm - # Whenever possible as it is the fastest backend - # is_aiter_found_and_supported() checks if aiter library is found - # and is aiter supported on the current platform - # aiter currently only is supported on gfx942 and gfx950 - # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py - if is_aiter_found_and_supported() and 90 < major * 10 + minor < 100: - if backend_name is None: - backend_name = "FLASH_ATTN" - else: - if backend_name == "FLASH_ATTN": - logger.warning( - "Flash Attention requires `aiter` library which is only supported " - "on gfx942 and gfx950. Falling back to TORCH_SDPA backend." - ) - backend_name = "TORCH_SDPA" - - if backend_name is not None: - backend_name_upper = backend_name.upper() - if backend_name_upper not in _BACKEND_CONFIG: - valid_backends = list(_BACKEND_CONFIG.keys()) - raise ValueError( - f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}" - ) - logger.info(f"Using attention backend '{backend_name_upper}' for diffusion") - return load_backend(backend_name_upper) - - # Default to SDPA - return SDPABackend + return _load_backend_cls(backend_cls_path) diff --git a/vllm_omni/diffusion/distributed/comm.py b/vllm_omni/diffusion/distributed/comm.py index ada03f4a1..f73b8442f 100644 --- a/vllm_omni/diffusion/distributed/comm.py +++ b/vllm_omni/diffusion/distributed/comm.py @@ -8,6 +8,8 @@ import torch import torch.distributed as dist from torch import Tensor +from vllm_omni.platforms import current_omni_platform + __all__ = ["all_to_all_4D", "all_to_all_5D", "SeqAllToAll4D", "SeqAllToAll5D", "RingComm"] @@ -48,7 +50,7 @@ def all_to_all_4D( if seq_world_size > 1: dist.all_to_all_single(output, input_t, group=group) if use_sync: - torch.cuda.synchronize() + current_omni_platform.synchronize() else: output = input_t # if scattering the seq-dim, transpose the heads back to the original dimension @@ -83,7 +85,7 @@ def all_to_all_4D( if seq_world_size > 1: dist.all_to_all_single(output, input_t, group=group) if use_sync: - torch.cuda.synchronize() + current_omni_platform.synchronize() else: output = input_t @@ -155,7 +157,7 @@ def all_to_all_5D( if seq_world_size > 1: dist.all_to_all_single(output, input_t, group=group) if use_sync: - torch.cuda.synchronize() + current_omni_platform.synchronize() else: output = input_t @@ -190,7 +192,7 @@ def all_to_all_5D( if seq_world_size > 1: dist.all_to_all_single(output, input_t, group=group) if use_sync: - torch.cuda.synchronize() + current_omni_platform.synchronize() else: output = input_t diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 8e33d6fb6..b722f61c0 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -9,18 +9,14 @@ from typing import Any import torch import torch.distributed -from torch.cuda import synchronize from torch.distributed import Backend, ProcessGroup from vllm.logger import init_logger from vllm_omni.diffusion import envs +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -if envs._is_npu(): - logger.info("torch.npu synchronize") - from torch.npu import synchronize - TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) @@ -124,7 +120,7 @@ class GroupCoordinator: assert self.cpu_group is not None assert self.device_group is not None - self.device = envs.get_device(local_rank) + self.device = current_omni_platform.get_torch_device(local_rank) @property def first_rank(self): @@ -618,7 +614,7 @@ class PipelineGroupCoordinator(GroupCoordinator): assert self.cpu_group is not None assert self.device_group is not None - self.device = envs.get_device(local_rank) + self.device = current_omni_platform.get_torch_device(local_rank) self.recv_buffer_set: bool = False self.recv_tasks_queue: list[tuple[str, int]] = [] @@ -768,7 +764,7 @@ class PipelineGroupCoordinator(GroupCoordinator): # To protect against race condition when using batch_isend_irecv(). # should take this out once the bug with batch_isend_irecv is resolved. - synchronize() + current_omni_platform.synchronize() ops = [] recv_prev_shape_tensor = None @@ -799,7 +795,7 @@ class PipelineGroupCoordinator(GroupCoordinator): for req in reqs: req.wait() - synchronize() + current_omni_platform.synchronize() recv_prev_shape = [0, 0, 0] if recv_prev_shape_tensor is not None: diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index b24951590..2f7c2f852 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm_omni.diffusion import envs +from vllm_omni.platforms import current_omni_platform from .group_coordinator import ( GroupCoordinator, @@ -42,14 +43,6 @@ from .group_coordinator import ( SequenceParallelGroupCoordinator, ) -if envs._is_npu(): - from torch.npu import device_count, set_device -elif envs._is_musa(): - from torch_musa.core.device import device_count, set_device -else: - from torch.cuda import device_count, set_device - - env_info = envs.PACKAGES_CHECKER.get_packages_info() HAS_FLASH_ATTN = env_info["has_flash_attn"] @@ -396,7 +389,7 @@ def init_distributed_environment( backend: str | None = None, ): if backend is None: - backend = envs.get_torch_distributed_backend() + backend = current_omni_platform.dist_backend logger.debug( "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s", world_size, @@ -416,7 +409,8 @@ def init_distributed_environment( world_size=world_size, rank=rank, ) - set_device(torch.distributed.get_rank() % device_count()) + device_id = torch.distributed.get_rank() % current_omni_platform.get_device_count() + current_omni_platform.set_device(current_omni_platform.get_torch_device(device_id)) # set the local rank # local_rank is not available in torch ProcessGroup, # see https://github.com/pytorch/pytorch/issues/122816 @@ -572,7 +566,7 @@ def initialize_model_parallel( backend: str | None = None, ) -> None: if backend is None: - backend = envs.get_torch_distributed_backend() + backend = current_omni_platform.dist_backend """ Initialize model parallel groups. @@ -632,7 +626,7 @@ def initialize_model_parallel( # FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch, # the pipefusion is not ready for npu yet - if envs._is_npu(): + if current_omni_platform.is_npu(): assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU" dit_parallel_size = ( diff --git a/vllm_omni/diffusion/distributed/utils.py b/vllm_omni/diffusion/distributed/utils.py index 2e2f0fa53..92d00ef4b 100644 --- a/vllm_omni/diffusion/distributed/utils.py +++ b/vllm_omni/diffusion/distributed/utils.py @@ -5,11 +5,10 @@ import os import torch -from vllm_omni.utils.platform_utils import detect_device_type +from vllm_omni.platforms import current_omni_platform def get_local_device() -> torch.device: """Return the torch device for the current rank based on detected device type.""" - device_type = detect_device_type() - local_rank = os.environ.get("LOCAL_RANK", 0) - return torch.device(f"{device_type}:{local_rank}") + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + return current_omni_platform.get_torch_device(local_rank) diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index f717f2ca0..be2b58ae8 100644 --- a/vllm_omni/diffusion/envs.py +++ b/vllm_omni/diffusion/envs.py @@ -5,20 +5,15 @@ import os from collections.abc import Callable from typing import TYPE_CHECKING, Any -import torch -from packaging import version from vllm.logger import init_logger -logger = init_logger(__name__) +from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: MASTER_ADDR: str = "" MASTER_PORT: int | None = None CUDA_HOME: str | None = None LOCAL_RANK: int = 0 - CUDA_VISIBLE_DEVICES: str | None = None - CUDA_VERSION: version.Version - TORCH_VERSION: version.Version environment_variables: dict[str, Callable[[], Any]] = { # ================== Runtime Env Vars ================== @@ -32,127 +27,14 @@ environment_variables: dict[str, Callable[[], Any]] = { # local rank of the process in the distributed setting, used to determine # the GPU device id "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), - # used to control the visible devices in the distributed setting - "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), } - -def _is_hip(): - has_rocm = torch.version.hip is not None - return has_rocm - - -def _is_cuda(): - has_cuda = torch.version.cuda is not None - return has_cuda - - -def _is_musa(): - try: - if hasattr(torch, "musa") and torch.musa.is_available(): - return True - except ModuleNotFoundError: - return False - - -def _is_mps(): - return torch.backends.mps.is_available() - - -def _is_npu(): - try: - if hasattr(torch, "npu") and torch.npu.is_available(): - return True - except ModuleNotFoundError: - return False - - -def get_device(local_rank: int) -> torch.device: - if _is_cuda() or _is_hip(): - return torch.device("cuda", local_rank) - elif _is_musa(): - return torch.device("musa", local_rank) - elif _is_mps(): - return torch.device("mps") - elif _is_npu(): - return torch.device("npu", local_rank) - else: - return torch.device("cpu") - - -def get_device_name() -> str: - if _is_cuda() or _is_hip(): - return "cuda" - elif _is_musa(): - return "musa" - elif _is_mps(): - return "mps" - elif _is_npu(): - return "npu" - else: - return "cpu" - - -def get_device_version(): - if _is_hip(): - hip_version = torch.version.hip - hip_version = hip_version.split("-")[0] - return hip_version - elif _is_cuda(): - return torch.version.cuda - elif _is_musa(): - return torch.version.musa - elif _is_mps(): - return None - elif _is_npu(): - return None - else: - raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available") - - -def get_torch_distributed_backend() -> str: - if _is_cuda() or _is_hip(): - return "nccl" - elif _is_musa(): - return "mccl" - elif _is_mps(): - return "gloo" - elif _is_npu(): - return "hccl" - else: - raise NotImplementedError("No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available") - - -variables: dict[str, Callable[[], Any]] = { - # ================== Other Vars ================== - # used in version checking - "CUDA_VERSION": lambda: version.parse(get_device_version() or "0.0"), - "TORCH_VERSION": lambda: version.parse(version.parse(torch.__version__).base_version), -} - - -def _setup_musa(environment_variables, variables): - musa = getattr(torch, "musa", None) - if musa is None: - return - try: - if musa.is_available(): - environment_variables["MUSA_HOME"] = lambda: os.environ.get("MUSA_HOME", None) - environment_variables["MUSA_VISIBLE_DEVICES"] = lambda: os.environ.get("MUSA_VISIBLE_DEVICES", None) - musa_ver = getattr(getattr(torch, "version", None), "musa", None) - if musa_ver: - variables["MUSA_VERSION"] = lambda: version.parse(musa_ver) - except Exception: - pass - - -try: - _setup_musa(environment_variables, variables) -except (AttributeError, ModuleNotFoundError): - pass +logger = init_logger(__name__) class PackagesEnvChecker: + """Singleton class for checking package availability.""" + _instance = None def __new__(cls): @@ -163,38 +45,39 @@ class PackagesEnvChecker: def initialize(self): packages_info = {} - packages_info["has_flash_attn"] = self.check_flash_attn(packages_info) + packages_info["has_flash_attn"] = self._check_flash_attn(packages_info) self.packages_info = packages_info - def check_flash_attn(self, packages_info): - if not torch.cuda.is_available(): + def _check_flash_attn(self, packages_info) -> bool: + """Check if flash attention is available and compatible.""" + platform = current_omni_platform + + # Flash attention requires CUDA-like platforms (CUDA or ROCm) + if not platform.is_cuda_alike(): return False - # Check if torch_npu is available - if _is_npu(): - logger.info("`flash_attn` is not ready on torch_npu for now") + # Check if devices are available + if platform.get_device_count() == 0: return False - if _is_musa(): - logger.info("Flash Attention library is not supported on MUSA for the moment.") - return False try: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - gpu_name = torch.cuda.get_device_name(device) + gpu_name = platform.get_device_name() + # Turing/Tesla/T4 GPUs don't support flash attention well if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name: return False - else: - from flash_attn import __version__ - if __version__ < "2.6.0": - raise ImportError("install flash_attn >= 2.6.0") - return True + from flash_attn import __version__ + + if __version__ < "2.6.0": + raise ImportError("install flash_attn >= 2.6.0") + return True except ImportError: if not packages_info.get("has_aiter", False): logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation') return False - def get_packages_info(self): + def get_packages_info(self) -> dict: + """Get the packages info dictionary.""" return self.packages_info @@ -205,8 +88,6 @@ def __getattr__(name): # lazy evaluation of environment variables if name in environment_variables: return environment_variables[name]() - if name in variables: - return variables[name]() raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 21a94acd3..a421a01ed 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -10,7 +10,7 @@ from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.scheduler import Scheduler -from vllm_omni.utils.platform_utils import get_diffusion_worker_class +from vllm_omni.diffusion.worker import WorkerProc logger = init_logger(__name__) @@ -77,9 +77,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): mp.set_start_method("spawn", force=True) processes = [] - # Get the appropriate worker class for current device - worker_proc = get_diffusion_worker_class() - # Launch all worker processes scheduler_pipe_readers = [] scheduler_pipe_writers = [] @@ -88,7 +85,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): reader, writer = mp.Pipe(duplex=False) scheduler_pipe_writers.append(writer) process = mp.Process( - target=worker_proc.worker_main, + target=WorkerProc.worker_main, args=( i, # rank od_config, diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 21f4a6fe1..8f2c2b755 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -335,10 +335,7 @@ class SequenceParallelSplitHook(ModelHook): 2. Creates an attention mask indicating valid vs padding positions 3. Stores the mask and padding info in ForwardContext """ - from vllm_omni.diffusion.attention.selector import ( - _BACKENDS_SUPPORT_ATTENTION_MASK, - get_attn_backend, - ) + from vllm_omni.diffusion.attention.selector import get_attn_backend from vllm_omni.diffusion.distributed.parallel_state import ( get_ring_parallel_world_size, get_sequence_parallel_rank, @@ -359,7 +356,7 @@ class SequenceParallelSplitHook(ModelHook): # Check backend compatibility attn_backend = get_attn_backend(-1) - if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK: + if not attn_backend.supports_attention_mask: raise ValueError( f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). " f"Cannot use {attn_backend.get_name()} which does not support attention_mask. " diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 835ee07c0..9a644e1d9 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -91,6 +91,14 @@ class AdaLayerNorm(CustomOp): return output, gate_result + def forward_xpu( + self, + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor = None, + ) -> torch.Tensor: + return self.forward_native(x, mod_params, index) + def forward_native( self, x: torch.Tensor, diff --git a/vllm_omni/diffusion/layers/custom_op.py b/vllm_omni/diffusion/layers/custom_op.py index 2f50073cc..321bcbf8a 100644 --- a/vllm_omni/diffusion/layers/custom_op.py +++ b/vllm_omni/diffusion/layers/custom_op.py @@ -3,7 +3,7 @@ from typing import Any import torch.nn as nn -from vllm_omni.utils.platform_utils import detect_device_type, is_npu, is_rocm +from vllm_omni.platforms import current_omni_platform class CustomOp(nn.Module): @@ -14,16 +14,17 @@ class CustomOp(nn.Module): def __init__(self) -> None: super().__init__() - self.is_cuda = detect_device_type() == "cuda" self._forward_method = self.dispatch_forward() def dispatch_forward(self) -> Callable: - if is_rocm(): + if current_omni_platform.is_rocm(): return self.forward_hip - elif self.is_cuda: + elif current_omni_platform.is_cuda(): return self.forward_cuda - elif is_npu(): + elif current_omni_platform.is_npu(): return self.forward_npu + elif current_omni_platform.is_xpu(): + return self.forward_xpu else: return self.forward_native @@ -44,6 +45,9 @@ class CustomOp(nn.Module): def forward_npu(self, *args, **kwargs): raise NotImplementedError + def forward_xpu(self, *args, **kwargs): + raise NotImplementedError + def forward_hip(self, *args, **kwargs): # By default, we assume that HIP ops are compatible with CUDA ops. return self.forward_cuda(*args, **kwargs) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 81e8c098f..461c25652 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -137,6 +137,14 @@ class RotaryEmbedding(CustomOp): else: return self.forward_native(x, cos, sin) + def forward_xpu( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x, cos, sin) + def forward_native( self, x: torch.Tensor, diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 6bb282a80..9d216363d 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -24,7 +24,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( get_sp_group, ) from vllm_omni.diffusion.forward_context import get_forward_context -from vllm_omni.utils.platform_utils import is_npu +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -610,7 +610,7 @@ class LongCatImageTransformer2DModel(nn.Module): ids = torch.cat((txt_ids, img_ids), dim=0) - if is_npu(): + if current_omni_platform.is_npu(): freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) else: diff --git a/vllm_omni/diffusion/worker/__init__.py b/vllm_omni/diffusion/worker/__init__.py index dfec4596b..8af028385 100644 --- a/vllm_omni/diffusion/worker/__init__.py +++ b/vllm_omni/diffusion/worker/__init__.py @@ -2,14 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Worker classes for diffusion models.""" -from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner -from vllm_omni.diffusion.worker.gpu_diffusion_worker import ( - GPUDiffusionWorker, +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.diffusion_worker import ( + DiffusionWorker, WorkerProc, ) __all__ = [ - "GPUDiffusionModelRunner", - "GPUDiffusionWorker", + "DiffusionModelRunner", + "DiffusionWorker", "WorkerProc", ] diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py similarity index 99% rename from vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py rename to vllm_omni/diffusion/worker/diffusion_model_runner.py index 0401655a8..7335162b0 100644 --- a/vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -34,7 +34,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec logger = init_logger(__name__) -class GPUDiffusionModelRunner: +class DiffusionModelRunner: """ Model runner that handles model loading and execution for diffusion models. diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py similarity index 94% rename from vllm_omni/diffusion/worker/gpu_diffusion_worker.py rename to vllm_omni/diffusion/worker/diffusion_worker.py index d8b962a58..0fad90a5b 100644 --- a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -4,7 +4,7 @@ Diffusion Worker for vLLM-Omni. Handles GPU infrastructure initialization and delegates model operations -to GPUDiffusionModelRunner. +to DiffusionModelRunner. """ import multiprocessing as mp @@ -31,13 +31,14 @@ from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.profiler import CurrentProfiler from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner +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 logger = init_logger(__name__) -class GPUDiffusionWorker: +class DiffusionWorker: """ A worker that manages GPU infrastructure and delegates to the model runner. @@ -47,7 +48,7 @@ class GPUDiffusionWorker: - Memory management (sleep/wake) All model-related operations (loading, compilation, execution) are - delegated to GPUDiffusionModelRunner. + delegated to DiffusionModelRunner. """ def __init__( @@ -61,7 +62,7 @@ class GPUDiffusionWorker: self.od_config = od_config self.device: torch.device | None = None self.vllm_config: VllmConfig | None = None - self.model_runner: GPUDiffusionModelRunner | None = None + self.model_runner: DiffusionModelRunner | None = None self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self.lora_manager: DiffusionLoRAManager | None = None self.init_device() @@ -79,8 +80,8 @@ class GPUDiffusionWorker: os.environ["WORLD_SIZE"] = str(world_size) # Setup device - self.device = torch.device(f"cuda:{rank}") - torch.cuda.set_device(self.device) + self.device = current_omni_platform.get_torch_device(rank) + current_omni_platform.set_device(self.device) # Create vllm_config for parallel configuration vllm_config = VllmConfig() @@ -105,7 +106,7 @@ class GPUDiffusionWorker: ) # Create model runner and load model - self.model_runner = GPUDiffusionModelRunner( + self.model_runner = DiffusionModelRunner( vllm_config=self.vllm_config, od_config=self.od_config, device=self.device, @@ -176,7 +177,7 @@ class GPUDiffusionWorker: """ from vllm.device_allocator.cumem import CuMemAllocator - free_bytes_before_sleep = torch.cuda.mem_get_info()[0] + 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: @@ -185,7 +186,9 @@ class GPUDiffusionWorker: allocator = CuMemAllocator.get_instance() allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) - free_bytes_after_sleep, total = torch.cuda.mem_get_info() + 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 assert freed_bytes >= 0, "Memory usage increased after sleeping." @@ -270,9 +273,9 @@ class WorkerProc: self.gpu_id = gpu_id self._running = True - def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> GPUDiffusionWorker: + def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> DiffusionWorker: """Create a worker instance. Override in subclasses for different worker types.""" - return GPUDiffusionWorker( + return DiffusionWorker( local_rank=gpu_id, rank=gpu_id, od_config=od_config, @@ -379,6 +382,9 @@ class WorkerProc: broadcast_handle, ) -> None: """Worker initialization and execution loops.""" + from vllm_omni.plugins import load_omni_general_plugins + + load_omni_general_plugins() worker_proc = WorkerProc( od_config, gpu_id=rank, diff --git a/vllm_omni/diffusion/worker/npu/__init__.py b/vllm_omni/diffusion/worker/npu/__init__.py deleted file mode 100644 index b2210b5e2..000000000 --- a/vllm_omni/diffusion/worker/npu/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NPU worker classes for diffusion models.""" - -from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorker, NPUWorkerProc - -__all__ = ["NPUWorker", "NPUWorkerProc"] diff --git a/vllm_omni/diffusion/worker/npu/npu_worker.py b/vllm_omni/diffusion/worker/npu/npu_worker.py deleted file mode 100644 index c4d2c60e5..000000000 --- a/vllm_omni/diffusion/worker/npu/npu_worker.py +++ /dev/null @@ -1,255 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import multiprocessing as mp -import os -import time -from collections.abc import Iterable -from contextlib import AbstractContextManager, nullcontext - -import torch -from vllm.config import LoadConfig, VllmConfig -from vllm.logger import init_logger -from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes - -from vllm_omni.diffusion.cache.selector import get_cache_backend -from vllm_omni.diffusion.data import ( - DiffusionOutput, - OmniDiffusionConfig, -) -from vllm_omni.diffusion.distributed.parallel_state import ( - destroy_distributed_env, - init_distributed_environment, - initialize_model_parallel, -) -from vllm_omni.diffusion.forward_context import set_forward_context -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc - -logger = init_logger(__name__) - - -class NPUWorker: - """ - A worker that executes the model on a single NPU. - Inherits from GPUWorker and overrides device-specific initialization. - """ - - def __init__( - self, - local_rank: int, - rank: int, - od_config: OmniDiffusionConfig, - ): - self.local_rank = local_rank - self.rank = rank - self.od_config = od_config - self.pipeline = None - self.device = None - self._sleep_saved_buffers: dict[str, torch.Tensor] = {} - self.init_device_and_model() - - def init_device_and_model(self) -> None: - """Initialize the NPU device and load the model.""" - world_size = self.od_config.num_gpus - rank = self.rank - # Set environment variables for distributed initialization - os.environ["MASTER_ADDR"] = "localhost" - os.environ["MASTER_PORT"] = str(self.od_config.master_port) - os.environ["LOCAL_RANK"] = str(self.local_rank) - os.environ["RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - - device = torch.device(f"npu:{rank}") - torch.npu.set_device(device) - - # hack - vllm_config = VllmConfig() - vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size - vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size - self.vllm_config = vllm_config - with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config): - init_distributed_environment(world_size=world_size, rank=rank) - logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") - parallel_config = self.od_config.parallel_config - initialize_model_parallel( - data_parallel_size=parallel_config.data_parallel_size, - cfg_parallel_size=parallel_config.cfg_parallel_size, - sequence_parallel_size=parallel_config.sequence_parallel_size, - ulysses_degree=parallel_config.ulysses_degree, - ring_degree=parallel_config.ring_degree, - tensor_parallel_size=parallel_config.tensor_parallel_size, - pipeline_parallel_size=parallel_config.pipeline_parallel_size, - ) - - load_config = LoadConfig() - model_loader = DiffusersPipelineLoader(load_config) - time_before_load = time.perf_counter() - with DeviceMemoryProfiler() as m: - self.pipeline = model_loader.load_model( - od_config=self.od_config, - load_device=f"npu:{rank}", - ) - time_after_load = time.perf_counter() - - logger.info( - "Model loading took %.4f GiB and %.6f seconds", - m.consumed_memory / GiB_bytes, - time_after_load - time_before_load, - ) - logger.info(f"Worker {self.rank}: Model loaded successfully.") - - # Setup cache backend based on type (both backends use enable()/reset() interface) - self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config) - - if self.cache_backend is not None: - self.cache_backend.enable(self.pipeline) - - 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.execute_model(requests, self.od_config) - - @torch.inference_mode() - def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput: - """ - Execute a forward pass. - """ - assert self.pipeline is not None - if not reqs or len(reqs) == 0: - raise ValueError("Cannot execute model with empty request list") - # TODO: dealing with first req for now - req = reqs[0] - - if req.generator is None and req.seed is not None: - req.generator = torch.Generator(device=self.device).manual_seed(req.seed) - - # Refresh cache context if needed - if ( - not getattr(req, "skip_cache_refresh", False) - and self.cache_backend is not None - and self.cache_backend.is_enabled() - ): - self.cache_backend.refresh(self.pipeline, req.num_inference_steps) - with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): - output = self.pipeline.forward(req) - return output - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - return self.pipeline.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. - """ - from vllm.device_allocator.cumem import CuMemAllocator - - free_bytes_before_sleep = torch.cuda.mem_get_info()[0] - - # Save the buffers before level 2 sleep - if level == 2: - model = self.pipeline - self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()} - - allocator = CuMemAllocator.get_instance() - allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) - free_bytes_after_sleep, total = torch.cuda.mem_get_info() - freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep - used_bytes = total - free_bytes_after_sleep - assert freed_bytes >= 0, "Memory usage increased after sleeping." - logger.info( - "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.", - freed_bytes / GiB_bytes, - used_bytes / GiB_bytes, - ) - return True - - 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. - """ - from vllm.device_allocator.cumem import CuMemAllocator - - allocator = CuMemAllocator.get_instance() - allocator.wake_up(tags) - - # Restore the buffers after level 2 sleep - if len(self._sleep_saved_buffers): - model = self.pipeline - for name, buffer in model.named_buffers(): - if name in self._sleep_saved_buffers: - buffer.data.copy_(self._sleep_saved_buffers[name].data) - self._sleep_saved_buffers = {} - return True - - def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: - if self.od_config.enable_sleep_mode: - from vllm.device_allocator.cumem import CuMemAllocator - - allocator = CuMemAllocator.get_instance() - if tag == "weights": - assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process." - return allocator.use_memory_pool(tag=tag) - else: - return nullcontext() - - def shutdown(self) -> None: - destroy_distributed_env() - - -class NPUWorkerProc(WorkerProc): - """Wrapper that runs one NPUWorker in a separate process.""" - - def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> NPUWorker: - """Create an NPUWorker instead of GPUWorker.""" - return NPUWorker( - local_rank=gpu_id, - rank=gpu_id, - od_config=od_config, - ) - - @staticmethod - def worker_main( - rank: int, - od_config: OmniDiffusionConfig, - pipe_writer: mp.connection.Connection, - broadcast_handle, - ) -> None: - """Worker initialization and execution loops.""" - - worker_proc = NPUWorkerProc( - od_config, - gpu_id=rank, - broadcast_handle=broadcast_handle, - ) - logger.info(f"Worker {rank}: Scheduler loop started.") - pipe_writer.send( - { - "status": "ready", - "result_handle": worker_proc.result_mq_handle if rank == 0 else None, - } - ) - worker_proc.worker_busy_loop() - logger.info(f"Worker {rank}: Shutdown complete.") diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index cbe2001d7..21bb761d2 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -7,6 +7,7 @@ from vllm.transformers_utils.config import get_hf_text_config from vllm.v1.engine.async_llm import AsyncEngineArgs from vllm_omni.config import OmniModelConfig +from vllm_omni.plugins import load_omni_general_plugins logger = init_logger(__name__) @@ -89,7 +90,11 @@ class OmniEngineArgs(EngineArgs): f"Config attribute '{hf_config_name}' not found in hf_config, " "falling back to default get_hf_text_config" ) - return get_hf_text_config(hf_config) + return get_hf_text_config(hf_config) + + def __post_init__(self) -> None: + load_omni_general_plugins() + super().__post_init__() def _ensure_omni_models_registered(self): if hasattr(self, "_omni_models_registered"): @@ -187,7 +192,11 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): f"Config attribute '{hf_config_name}' not found in hf_config, " "falling back to default get_hf_text_config" ) - return get_hf_text_config(hf_config) + return get_hf_text_config(hf_config) + + def __post_init__(self) -> None: + load_omni_general_plugins() + super().__post_init__() def _ensure_omni_models_registered(self): if hasattr(self, "_omni_models_registered"): diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 0a5ee55be..2adb1c342 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -51,11 +51,27 @@ from vllm_omni.entrypoints.stage_utils import ( ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils import detect_device_type logger = init_logger(__name__) +def _resolve_worker_cls(engine_args: dict[str, Any]) -> None: + worker_type = engine_args.pop("worker_type", None) + if not worker_type: + return + if engine_args.get("worker_cls"): + return + from vllm_omni.platforms import current_omni_platform + + worker_type = str(worker_type).lower() + if worker_type == "ar": + engine_args["worker_cls"] = current_omni_platform.get_omni_ar_worker_cls() + elif worker_type == "generation": + engine_args["worker_cls"] = current_omni_platform.get_omni_generation_worker_cls() + else: + raise ValueError(f"Unknown worker_type: {worker_type}") + + def _build_od_config(engine_args: dict[str, Any], model: str) -> dict[str, Any]: """Build OmniDiffusionConfig kwargs from engine args.""" od_config = engine_args.get("od_config", {}) @@ -483,6 +499,9 @@ def _stage_worker( import os as _os import time as _time + from vllm_omni.plugins import load_omni_general_plugins + + load_omni_general_plugins() # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker / # GPUARModelRunner) are spawned with a fork-safe method. # Mooncake / gRPC / RDMA and CUDA/NCCL can deadlock under fork-with-threads. @@ -503,6 +522,9 @@ def _stage_worker( connectors_config = stage_payload.get("connectors_config", {}) stage_type: Literal["llm", "diffusion"] = stage_payload.get("stage_type", "llm") + if stage_type != "diffusion": + _resolve_worker_cls(engine_args) + # Aggregates for running average _agg_total_tokens = 0 _agg_total_gen_time_ms = 0.0 @@ -512,7 +534,9 @@ def _stage_worker( # Device mapping device_type = None try: - device_type = detect_device_type() + from vllm_omni.platforms import current_omni_platform + + device_type = current_omni_platform.device_type set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type) except Exception as e: logger.warning("Device setup failed: %s", e) @@ -521,139 +545,137 @@ def _stage_worker( # when multiple instances start simultaneously # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage lock_files = [] - if device_type == "cuda": - try: - import torch + try: + # Get all parallel sizes from engine_args or parallel_config (defaults to 1) + if "parallel_config" in engine_args: + parallel_config = engine_args["parallel_config"] + tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) + pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) + data_parallel_size = parallel_config.get("data_parallel_size", 1) + prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) + sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) + cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) + else: + tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) + pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) + data_parallel_size = engine_args.get("data_parallel_size", 1) + prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) + sequence_parallel_size = 1 # not use in omni model + cfg_parallel_size = 1 # not used in omni model - if torch.cuda.is_available(): - # Get all parallel sizes from engine_args or parallel_config (defaults to 1) - if "parallel_config" in engine_args: - parallel_config = engine_args["parallel_config"] - tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) - pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) - data_parallel_size = parallel_config.get("data_parallel_size", 1) - prefill_context_parallel_size = 1 # not used for diffusion - sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) - cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) - else: - tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) - pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) - data_parallel_size = engine_args.get("data_parallel_size", 1) - prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) - sequence_parallel_size = 1 # not use in omni model - cfg_parallel_size = 1 # not used in omni model + # Calculate total number of devices needed for this stage + # For a single stage worker: + # - TP: splits model across devices (always needed) + # - PP: splits layers across pipeline stages, but each stage uses TP devices + # - DP: replicates model, but each replica uses TP devices + # - PCP: context parallelism, typically uses TP devices + # - SP: sequence parallelism, typically uses TP devices + # - CFG: Classifier-Free Guidance parallelism for diffusion models + # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size + # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) + num_devices_per_stage = ( + tensor_parallel_size + * pipeline_parallel_size + * data_parallel_size + * prefill_context_parallel_size + * sequence_parallel_size + * cfg_parallel_size + ) - # Calculate total number of devices needed for this stage - # For a single stage worker: - # - TP: splits model across GPUs (always needed) - # - PP: splits layers across pipelinestages, but each stage uses TP devices - # - DP: replicates model, but each replica uses TP devices - # - PCP: context parallelism, typically uses TP devices - # - SP: sequence parallelism, typically uses TP devices - # - CFG: Classifier-Free Guidance parallelism for diffusion models - # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size - # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) - num_devices_per_stage = ( - tensor_parallel_size - * pipeline_parallel_size - * data_parallel_size - * prefill_context_parallel_size - * sequence_parallel_size - * cfg_parallel_size - ) + # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES) + # After set_stage_devices, this env var is set to physical device(s) + device_control_env = current_omni_platform.device_control_env_var + visible_devices_str = _os.environ.get(device_control_env) + physical_devices = [] - # Get physical device IDs from CUDA_VISIBLE_DEVICES - # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s) - cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES") - physical_devices = [] + if visible_devices_str: + try: + physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] + except (ValueError, IndexError): + pass - if cuda_visible_devices: + if not physical_devices: + # Fallback: use logical device count if device control env var not set + num_devices = current_omni_platform.get_device_count() + physical_devices = list(range(num_devices)) + + # Determine which devices will be used (min of devices per stage and available devices) + num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) + devices_to_lock = physical_devices[:num_devices_to_lock] + + # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) + devices_to_lock = sorted(devices_to_lock) + + logger.debug( + "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", + tensor_parallel_size, + pipeline_parallel_size, + data_parallel_size, + prefill_context_parallel_size, + sequence_parallel_size, + cfg_parallel_size, + num_devices_to_lock, + devices_to_lock, + ) + + # Acquire exclusive locks for all devices using fcntl.flock + # Locks are automatically released when process dies + wait_start = _time.time() + acquired_lock_fds = [] # Store file descriptors to keep locks alive + + for device_id in devices_to_lock: + lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" + lock_acquired = False + + while not lock_acquired: + try: + # Open or create the lock file + lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) + + # Try to acquire exclusive lock (non-blocking first) try: - physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()] - except (ValueError, IndexError): - pass + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + # Successfully acquired lock - write PID + _os.ftruncate(lock_fd, 0) # Clear file + _os.write(lock_fd, f"{_os.getpid()}\n".encode()) + _os.fsync(lock_fd) # Ensure written to disk + lock_acquired = True + acquired_lock_fds.append(lock_fd) + logger.debug("Acquired exclusive lock for device %s", device_id) + except BlockingIOError: + # Lock is held by another process + _os.close(lock_fd) - if not physical_devices: - # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set - num_devices = torch.cuda.device_count() - physical_devices = list(range(num_devices)) - - # Determine which devices will be used (min of devices per stage and available devices) - num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) - devices_to_lock = physical_devices[:num_devices_to_lock] - - # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) - devices_to_lock = sorted(devices_to_lock) - - logger.debug( - "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d; will lock %d devices: %s", - tensor_parallel_size, - pipeline_parallel_size, - data_parallel_size, - prefill_context_parallel_size, - sequence_parallel_size, - num_devices_to_lock, - devices_to_lock, - ) - - # Acquire exclusive locks for all devices using fcntl.flock - # Locks are automatically released when process dies - wait_start = _time.time() - acquired_lock_fds = [] # Store file descriptors to keep locks alive - - for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" - lock_acquired = False - - while not lock_acquired: - try: - # Open or create the lock file - lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) - - # Try to acquire exclusive lock (non-blocking first) - try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Successfully acquired lock - write PID - _os.ftruncate(lock_fd, 0) # Clear file - _os.write(lock_fd, f"{_os.getpid()}\n".encode()) - _os.fsync(lock_fd) # Ensure written to disk - lock_acquired = True - acquired_lock_fds.append(lock_fd) - logger.debug("Acquired exclusive lock for device %s", device_id) - except BlockingIOError: - # Lock is held by another process - _os.close(lock_fd) - - # Check if we've been waiting too long - if _time.time() - wait_start > stage_init_timeout: - logger.warning( - "Timeout waiting for device %s initialization lock, proceeding anyway", - device_id, - ) - break - - # Wait a bit before retrying - _time.sleep(0.1) - except OSError as e: - # Other error - log and continue without lock - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", + # Check if we've been waiting too long + if _time.time() - wait_start > stage_init_timeout: + logger.warning( + "Timeout waiting for device %s initialization lock, proceeding anyway", device_id, - e, ) - try: - _os.close(lock_fd) - except (OSError, NameError): - pass break - lock_files = acquired_lock_fds - except Exception as e: - logger.debug( - "[Stage-%s] Failed to set up sequential initialization lock: %s", - stage_id, - e, - ) + # Wait a bit before retrying + _time.sleep(0.1) + except OSError as e: + # Other error - log and continue without lock + logger.debug( + "Failed to acquire lock for device %s: %s, continuing anyway", + device_id, + e, + ) + try: + _os.close(lock_fd) + except (OSError, NameError): + pass + break + + lock_files = acquired_lock_fds + except Exception as e: + logger.debug( + "[Stage-%s] Failed to set up sequential initialization lock: %s", + stage_id, + e, + ) # Init engine based on stage_type logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())) if engine_args.get("async_chunk", False): @@ -1001,6 +1023,9 @@ async def _stage_worker_async( import os as _os import time as _time + from vllm_omni.plugins import load_omni_general_plugins + + load_omni_general_plugins() # IMPORTANT: Ensure vLLM's internal multiprocessing workers (e.g., GPUARWorker / # GPUARModelRunner) are spawned with a fork-safe method. if _os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn": @@ -1018,6 +1043,9 @@ async def _stage_worker_async( connectors_config = stage_payload.get("connectors_config", {}) stage_type = stage_payload.get("stage_type", "llm") + if stage_type != "diffusion": + _resolve_worker_cls(engine_args) + in_q = omni_stage._in_q out_q = omni_stage._out_q @@ -1031,9 +1059,9 @@ async def _stage_worker_async( # Device mapping device_type = None try: - from vllm_omni.utils import detect_device_type + from vllm_omni.platforms import current_omni_platform - device_type = detect_device_type() + device_type = current_omni_platform.device_type set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type) except Exception as e: logger.warning("Device setup failed: %s", e) @@ -1051,120 +1079,136 @@ async def _stage_worker_async( # Sequential initialization on the same device to avoid memory calculation errors # when multiple instances start simultaneously - # For TP, we need to lock ALL devices that will be used by this stage + # For TP/PP/DP/PCP, we need to lock ALL devices that will be used by this stage lock_files = [] - if device_type == "cuda": - try: - import torch + try: + # Get all parallel sizes from engine_args or parallel_config (defaults to 1) + if "parallel_config" in engine_args: + parallel_config = engine_args["parallel_config"] + tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) + pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) + data_parallel_size = parallel_config.get("data_parallel_size", 1) + prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) + sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) + cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) + else: + tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) + pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) + data_parallel_size = engine_args.get("data_parallel_size", 1) + prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) + sequence_parallel_size = 1 # not use in omni model + cfg_parallel_size = 1 # not used in omni model - if torch.cuda.is_available(): - # Get all parallel sizes from engine_args (defaults to 1) - tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) - pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) - data_parallel_size = engine_args.get("data_parallel_size", 1) - prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) + # Calculate total number of devices needed for this stage + # For a single stage worker: + # - TP: splits model across devices (always needed) + # - PP: splits layers across pipeline stages, but each stage uses TP devices + # - DP: replicates model, but each replica uses TP devices + # - PCP: context parallelism, typically uses TP devices + # - SP: sequence parallelism, typically uses TP devices + # - CFG: Classifier-Free Guidance parallelism for diffusion models + # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size + # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) + num_devices_per_stage = ( + tensor_parallel_size + * pipeline_parallel_size + * data_parallel_size + * prefill_context_parallel_size + * sequence_parallel_size + * cfg_parallel_size + ) - # Calculate total number of devices needed for this stage - # For a single stage worker in omni: - # - TP: splits model across GPUs (always needed) - # - PP: splits layers across stages, but each stage uses TP devices - # - DP: replicates model, but each replica uses TP devices - # - PCP: context parallelism, typically uses TP devices - # The number of devices per stage is determined by TP * PP * DP * PCP size - # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) - num_devices_per_stage = ( - tensor_parallel_size * pipeline_parallel_size * data_parallel_size * prefill_context_parallel_size - ) + # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES) + # After set_stage_devices, this env var is set to physical device(s) + device_control_env = current_omni_platform.device_control_env_var + visible_devices_str = _os.environ.get(device_control_env) + physical_devices = [] - # Get physical device IDs from CUDA_VISIBLE_DEVICES - # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s) - cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES") - physical_devices = [] + if visible_devices_str: + try: + physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] + except (ValueError, IndexError): + pass - if cuda_visible_devices: + if not physical_devices: + # Fallback: use logical device count if device control env var not set + num_devices = current_omni_platform.get_device_count() + physical_devices = list(range(num_devices)) + + # Determine which devices will be used (min of devices per stage and available devices) + num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) + devices_to_lock = physical_devices[:num_devices_to_lock] + + # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) + devices_to_lock = sorted(devices_to_lock) + + logger.debug( + "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", + tensor_parallel_size, + pipeline_parallel_size, + data_parallel_size, + prefill_context_parallel_size, + sequence_parallel_size, + cfg_parallel_size, + num_devices_to_lock, + devices_to_lock, + ) + + # Acquire exclusive locks for all devices using fcntl.flock + # Locks are automatically released when process dies + wait_start = _time.time() + acquired_lock_fds = [] # Store file descriptors to keep locks alive + + for device_id in devices_to_lock: + lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" + lock_acquired = False + + while not lock_acquired: + try: + # Open or create the lock file + lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) + + # Try to acquire exclusive lock (non-blocking first) try: - physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()] - except (ValueError, IndexError): - pass + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + # Successfully acquired lock - write PID + _os.ftruncate(lock_fd, 0) # Clear file + _os.write(lock_fd, f"{_os.getpid()}\n".encode()) + _os.fsync(lock_fd) # Ensure written to disk + lock_acquired = True + acquired_lock_fds.append(lock_fd) + logger.debug("Acquired exclusive lock for device %s", device_id) + except BlockingIOError: + # Lock is held by another process + _os.close(lock_fd) - if not physical_devices: - # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set - num_devices = torch.cuda.device_count() - physical_devices = list(range(num_devices)) - - # Determine which devices will be used (min of devices per stage and available devices) - num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) - devices_to_lock = physical_devices[:num_devices_to_lock] - - # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) - devices_to_lock = sorted(devices_to_lock) - - logger.debug( - "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s", - tensor_parallel_size, - pipeline_parallel_size, - data_parallel_size, - prefill_context_parallel_size, - num_devices_to_lock, - devices_to_lock, - ) - - # Acquire exclusive locks for all devices using fcntl.flock - # Locks are automatically released when process dies - wait_start = _time.time() - acquired_lock_fds = [] # Store file descriptors to keep locks alive - - for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" - lock_acquired = False - - while not lock_acquired: - try: - # Open or create the lock file - lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) - - # Try to acquire exclusive lock (non-blocking first) - try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Successfully acquired lock - write PID - _os.ftruncate(lock_fd, 0) # Clear file - _os.write(lock_fd, f"{_os.getpid()}\n".encode()) - _os.fsync(lock_fd) # Ensure written to disk - lock_acquired = True - acquired_lock_fds.append(lock_fd) - logger.debug("Acquired exclusive lock for device %s", device_id) - except BlockingIOError: - # Lock is held by another process - _os.close(lock_fd) - - # Check if we've been waiting too long - if _time.time() - wait_start > stage_init_timeout: - logger.warning( - "Timeout waiting for device %s initialization lock, " - "proceeding anyway with timeout %s", - device_id, - stage_init_timeout, - ) - break - - # Wait a bit before retrying - _time.sleep(0.1) - except OSError as e: - # Other error - log and continue without lock - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", + # Check if we've been waiting too long + if _time.time() - wait_start > stage_init_timeout: + logger.warning( + "Timeout waiting for device %s initialization lock, proceeding anyway with timeout %s", device_id, - e, + stage_init_timeout, ) - try: - _os.close(lock_fd) - except (OSError, NameError): - pass break - lock_files = acquired_lock_fds - except Exception as e: - logger.debug("Failed to set up sequential initialization lock: %s", e) + # Wait a bit before retrying + _time.sleep(0.1) + except OSError as e: + # Other error - log and continue without lock + logger.debug( + "Failed to acquire lock for device %s: %s, continuing anyway", + device_id, + e, + ) + try: + _os.close(lock_fd) + except (OSError, NameError): + pass + break + + lock_files = acquired_lock_fds + except Exception as e: + logger.debug("Failed to set up sequential initialization lock: %s", e) # Init engine based on stage_type logger.debug( diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index da837ce9a..74ad42f04 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -57,12 +57,12 @@ def set_stage_devices( - CUDA: Sets CUDA_VISIBLE_DEVICES and calls torch.cuda.set_device() - NPU: Sets ASCEND_RT_VISIBLE_DEVICES and calls torch.npu.set_device() """ - from vllm_omni.utils import detect_device_type, get_device_control_env_var + from vllm_omni.platforms import current_omni_platform if device_type is None: - device_type = detect_device_type() + device_type = current_omni_platform.device_type - env_var = get_device_control_env_var() + env_var = current_omni_platform.device_control_env_var try: selected_physical: int | None = None diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index f1ed4d843..4ebff3290 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -10,7 +10,7 @@ from vllm.transformers_utils.config import get_config, get_hf_file_to_dict from vllm.transformers_utils.repo_utils import file_or_path_exists from vllm_omni.entrypoints.stage_utils import _to_dict -from vllm_omni.utils import detect_device_type, is_rocm +from vllm_omni.platforms import current_omni_platform # Get the project root directory (2 levels up from this file) PROJECT_ROOT = Path(__file__).parent.parent.parent @@ -164,16 +164,12 @@ def resolve_model_config_path(model: str) -> str: f"Model is not in standard transformers format and does not have model_index.json. " f"Please ensure the model has proper configuration files with 'model_type' field" ) - device_type = detect_device_type() - # Try device-specific config first - if device_type != "cuda" or is_rocm(): - device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml" - if is_rocm(): - device_config_file = f"vllm_omni/model_executor/stage_configs/rocm/{model_type}.yaml" - device_config_path = PROJECT_ROOT / device_config_file - if os.path.exists(device_config_path): - return str(device_config_path) + default_config_path = current_omni_platform.get_default_stage_config_path() + model_type_str = f"{model_type}.yaml" + complete_config_path = PROJECT_ROOT / default_config_path / model_type_str + if os.path.exists(complete_config_path): + return str(complete_config_path) # Fall back to default config stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml" diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 45a1447b3..f4afe675a 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -36,7 +36,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import download_weights_ from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision -from vllm_omni.utils.platform_utils import is_npu +from vllm_omni.platforms import current_omni_platform TALKER_CODEC_EOS_TOKEN_ID = 8294 TALKER_CODEC_BOS_TOKEN_ID = 8293 @@ -247,7 +247,7 @@ class Qwen2_5OmniForConditionalGeneration( if inputs_embeds is not None and inputs_embeds.device != thinker_dev: inputs_embeds = inputs_embeds.to(thinker_dev) - if is_npu(): + if current_omni_platform.is_npu(): # TODO: remove this hack when NPU supports batched inputs properly thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py index c60015c0e..9d176ac6c 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py @@ -35,7 +35,7 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames -from vllm_omni.utils.platform_utils import is_npu +from vllm_omni.platforms import current_omni_platform # Provide a no-op auto_docstring decorator to satisfy annotations if missing @@ -732,7 +732,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> beta = 0.0 # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument - if is_npu(): + if current_omni_platform.is_npu(): kaiser_window = torch.kaiser_window( kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu" ).to("npu") @@ -796,7 +796,7 @@ class UpSample1d(nn.Module): def forward(self, hidden_states): channels = hidden_states.shape[1] - if is_npu(): + if current_omni_platform.is_npu(): # TODO: When F.pad supports replicate mode on NPU, remove this branch input_dtype = hidden_states.dtype # F.pad in NPU doesn't support BF16 when mode is replicate. @@ -843,7 +843,7 @@ class DownSample1d(nn.Module): def forward(self, hidden_states): channels = hidden_states.shape[1] - if is_npu(): + if current_omni_platform.is_npu(): input_dtype = hidden_states.dtype # F.pad in NPU doesn't support BF16 when mode is replicate. # To ensure the accuracy, manually pad the input tensor. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index a42bd8e76..37753ee47 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -35,7 +35,7 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape -from vllm_omni.utils.platform_utils import is_npu +from vllm_omni.platforms import current_omni_platform # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -269,7 +269,7 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 1: Thinker ========== if self.model_stage == "thinker": thinker_dev = self._module_device(self.thinker) - if is_npu(): + if current_omni_platform.is_npu(): # Normalize to batched inputs if needed _added_batch_dim = False if input_ids is not None and input_ids.ndim == 1: @@ -308,7 +308,7 @@ class Qwen3OmniMoeForConditionalGeneration( "capture_layer_indices": [0, int(accept_layer)], "return_hidden_states": True, } - if is_npu(): + if current_omni_platform.is_npu(): # TODO: remove this hack when NPU supports batched inputs properly thinker_input_ids = input_ids[0] if input_ids is not None and _added_batch_dim else input_ids thinker_inputs_embeds = ( diff --git a/vllm_omni/model_executor/stage_configs/bagel.yaml b/vllm_omni/model_executor/stage_configs/bagel.yaml index 96bdc01b6..93d2b611a 100644 --- a/vllm_omni/model_executor/stage_configs/bagel.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel.yaml @@ -9,7 +9,7 @@ stage_args: engine_args: model_stage: thinker model_arch: BagelForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.4 enforce_eager: true diff --git a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml index 678d5775f..86c0330a6 100644 --- a/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml @@ -9,7 +9,7 @@ stage_args: engine_args: model_stage: thinker model_arch: BagelForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.4 enforce_eager: true diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index e3d740ad5..e6ed97660 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -11,7 +11,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true # Now we only support eager mode @@ -40,7 +40,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true @@ -69,7 +69,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml index aae194d37..68f5817a7 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml @@ -11,7 +11,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true # Now we only support eager mode @@ -41,7 +41,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true @@ -73,7 +73,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.3 enforce_eager: true 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 9c422d76e..d0380a39d 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -14,7 +14,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: true @@ -45,7 +45,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: true @@ -76,7 +76,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true 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 55d455f56..42748960e 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 @@ -14,7 +14,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -46,7 +46,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: false @@ -77,7 +77,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml index 4abd3835a..69102df6f 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml @@ -13,7 +13,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: true @@ -46,7 +46,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: true @@ -82,7 +82,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml index 40799fe7f..d408dbab9 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -7,7 +7,7 @@ stage_args: engine_args: model_stage: qwen3_tts model_arch: Qwen3TTSForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/platforms/__init__.py b/vllm_omni/platforms/__init__.py new file mode 100644 index 000000000..b76c59a67 --- /dev/null +++ b/vllm_omni/platforms/__init__.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +import traceback +from itertools import chain +from typing import TYPE_CHECKING + +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.utils.torch_utils import supports_xccl + +from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum +from vllm_omni.plugins import ( + OMNI_PLATFORM_PLUGINS_GROUP, + load_omni_plugins_by_group, +) + +logger = logging.getLogger(__name__) + + +def cuda_omni_platform_plugin() -> str | None: + """Check if CUDA OmniPlatform should be activated.""" + is_cuda = False + logger.debug("Checking if CUDA OmniPlatform is available.") + try: + from vllm.utils.import_utils import import_pynvml + + pynvml = import_pynvml() + pynvml.nvmlInit() + try: + if pynvml.nvmlDeviceGetCount() > 0: + is_cuda = True + logger.debug("Confirmed CUDA OmniPlatform is available.") + else: + logger.debug("CUDA OmniPlatform is not available because no GPU is found.") + finally: + pynvml.nvmlShutdown() + except Exception as e: + logger.debug("CUDA OmniPlatform is not available because: %s", str(e)) + + return "vllm_omni.platforms.cuda.platform.CudaOmniPlatform" if is_cuda else None + + +def rocm_omni_platform_plugin() -> str | None: + """Check if ROCm OmniPlatform should be activated.""" + is_rocm = False + logger.debug("Checking if ROCm OmniPlatform is available.") + try: + import amdsmi + + amdsmi.amdsmi_init() + try: + if len(amdsmi.amdsmi_get_processor_handles()) > 0: + is_rocm = True + logger.debug("Confirmed ROCm OmniPlatform is available.") + else: + logger.debug("ROCm OmniPlatform is not available because no GPU is found.") + finally: + amdsmi.amdsmi_shut_down() + except Exception as e: + logger.debug("ROCm OmniPlatform is not available because: %s", str(e)) + + return "vllm_omni.platforms.rocm.platform.RocmOmniPlatform" if is_rocm else None + + +def npu_omni_platform_plugin() -> str | None: + """Check if NPU OmniPlatform should be activated.""" + is_npu = False + logger.debug("Checking if NPU OmniPlatform is available.") + try: + import torch + + if hasattr(torch, "npu") and torch.npu.is_available(): + is_npu = True + logger.debug("Confirmed NPU OmniPlatform is available.") + except Exception as e: + logger.debug("NPU OmniPlatform is not available because: %s", str(e)) + + return "vllm_omni.platforms.npu.platform.NPUOmniPlatform" if is_npu else None + + +def xpu_omni_platform_plugin() -> str | None: + """Check if XPU OmniPlatform should be activated.""" + is_xpu = False + logger.debug("Checking if XPU OmniPlatform is available.") + try: + # installed IPEX if the machine has XPUs. + import intel_extension_for_pytorch # noqa: F401 + import torch + + if supports_xccl(): + dist_backend = "xccl" + else: + dist_backend = "ccl" + import oneccl_bindings_for_pytorch # noqa: F401 + + if hasattr(torch, "xpu") and torch.xpu.is_available(): + is_xpu = True + from vllm_omni.platforms.xpu import XPUOmniPlatform + + XPUOmniPlatform.dist_backend = dist_backend + logger.debug("Confirmed %s backend is available.", XPUOmniPlatform.dist_backend) + logger.debug("Confirmed XPU platform is available.") + except Exception as e: + logger.debug("XPU omni platform is not available because: %s", str(e)) + + return "vllm_omni.platforms.xpu.platform.XPUOmniPlatform" if is_xpu else None + + +builtin_omni_platform_plugins = { + "cuda": cuda_omni_platform_plugin, + "rocm": rocm_omni_platform_plugin, + "npu": npu_omni_platform_plugin, + "xpu": xpu_omni_platform_plugin, +} + + +def resolve_current_omni_platform_cls_qualname() -> str: + """Resolve the current OmniPlatform class qualified name.""" + platform_plugins = load_omni_plugins_by_group(OMNI_PLATFORM_PLUGINS_GROUP) + + activated_plugins = [] + + for name, func in chain(builtin_omni_platform_plugins.items(), platform_plugins.items()): + try: + assert callable(func) + platform_cls_qualname = func() + if platform_cls_qualname is not None: + activated_plugins.append(name) + except Exception: + pass + + activated_builtin_plugins = list(set(activated_plugins) & set(builtin_omni_platform_plugins.keys())) + activated_oot_plugins = list(set(activated_plugins) & set(platform_plugins.keys())) + + if len(activated_oot_plugins) >= 2: + raise RuntimeError(f"Only one OmniPlatform plugin can be activated, but got: {activated_oot_plugins}") + elif len(activated_oot_plugins) == 1: + platform_cls_qualname = platform_plugins[activated_oot_plugins[0]]() + logger.info("OmniPlatform plugin %s is activated", activated_oot_plugins[0]) + elif len(activated_builtin_plugins) >= 2: + raise RuntimeError(f"Only one OmniPlatform plugin can be activated, but got: {activated_builtin_plugins}") + elif len(activated_builtin_plugins) == 1: + platform_cls_qualname = builtin_omni_platform_plugins[activated_builtin_plugins[0]]() + logger.debug("Automatically detected OmniPlatform %s.", activated_builtin_plugins[0]) + else: + platform_cls_qualname = "vllm_omni.platforms.interface.UnspecifiedOmniPlatform" + logger.debug("No platform detected, vLLM-Omni is running on UnspecifiedOmniPlatform") + + return platform_cls_qualname + + +_current_omni_platform = None +_init_trace: str = "" + +if TYPE_CHECKING: + current_omni_platform: OmniPlatform + + +def __getattr__(name: str): + if name == "current_omni_platform": + # Lazy init current_omni_platform + global _current_omni_platform + if _current_omni_platform is None: + platform_cls_qualname = resolve_current_omni_platform_cls_qualname() + _current_omni_platform = resolve_obj_by_qualname(platform_cls_qualname)() + global _init_trace + _init_trace = "".join(traceback.format_stack()) + return _current_omni_platform + elif name in globals(): + return globals()[name] + else: + raise AttributeError(f"No attribute named '{name}' exists in {__name__}.") + + +def __setattr__(name: str, value): # noqa: N807 + if name == "current_omni_platform": + global _current_omni_platform + _current_omni_platform = value + elif name in globals(): + globals()[name] = value + else: + raise AttributeError(f"No attribute named '{name}' exists in {__name__}.") + + +__all__ = [ + "OmniPlatform", + "OmniPlatformEnum", + "current_omni_platform", + "_init_trace", +] diff --git a/vllm_omni/platforms/cuda/__init__.py b/vllm_omni/platforms/cuda/__init__.py new file mode 100644 index 000000000..91ef864be --- /dev/null +++ b/vllm_omni/platforms/cuda/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + +__all__ = ["CudaOmniPlatform"] diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py new file mode 100644 index 000000000..9c65877be --- /dev/null +++ b/vllm_omni/platforms/cuda/platform.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.logger import init_logger +from vllm.platforms.cuda import CudaPlatformBase +from vllm.platforms.interface import DeviceCapability + +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum + +logger = init_logger(__name__) + + +class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): + """CUDA/GPU implementation of OmniPlatform (default). + + Inherits all CUDA-specific implementations from vLLM's CudaPlatform, + and adds Omni-specific interfaces from OmniPlatform. + """ + + _omni_enum = OmniPlatformEnum.CUDA + + @classmethod + def get_omni_ar_worker_cls(cls) -> str: + return "vllm_omni.worker.gpu_ar_worker.GPUARWorker" + + @classmethod + def get_omni_generation_worker_cls(cls) -> str: + return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker" + + @classmethod + def get_default_stage_config_path(cls) -> str: + return "vllm_omni/model_executor/stage_configs" + + @classmethod + def get_diffusion_attn_backend_cls( + cls, + selected_backend: str | None, + head_size: int, + ) -> str: + # Check compute capability for Flash Attention support + # Flash Attention requires compute capability >= 8.0 and < 10.0 + compute_capability = cls.get_device_capability() + flash_attn_supported = False + if compute_capability is not None: + major, minor = compute_capability + capability = major * 10 + minor + flash_attn_supported = 80 <= capability < 100 + + if selected_backend is not None: + backend_upper = selected_backend.upper() + if backend_upper == "FLASH_ATTN" and not flash_attn_supported: + logger.warning( + "Flash Attention requires GPU with compute capability >= 8.0 " + "and < 10.0. Falling back to TORCH_SDPA backend." + ) + logger.info("Defaulting to diffusion attention backend SDPA") + return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + backend = DiffusionAttentionBackendEnum[backend_upper] + logger.info("Using diffusion attention backend '%s'", backend_upper) + return backend.get_path() + + if flash_attn_supported: + logger.info("Defaulting to diffusion attention backend FLASH_ATTN") + return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + + logger.info("Defaulting to diffusion attention backend SDPA") + return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + if local_rank is None: + return torch.device("cuda") + return torch.device("cuda", local_rank) + + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + + @classmethod + def get_device_count(cls) -> int: + return torch.cuda.device_count() + + @classmethod + def get_device_version(cls) -> str | None: + return torch.version.cuda + + @classmethod + def synchronize(cls) -> None: + torch.cuda.synchronize() + + @classmethod + def get_free_memory(cls, device: torch.device | None = None) -> int: + free, _ = torch.cuda.mem_get_info(device) + return free + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return torch.cuda.get_device_name(device_id) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py new file mode 100644 index 000000000..2aac4b7a1 --- /dev/null +++ b/vllm_omni/platforms/interface.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import Enum + +import torch +from vllm.platforms import Platform + + +class OmniPlatformEnum(Enum): + """Enum for supported Omni platforms.""" + + CUDA = "cuda" + ROCM = "rocm" + NPU = "npu" + XPU = "xpu" + UNSPECIFIED = "unspecified" + + +class OmniPlatform(Platform): + """ + Abstract base class for vllm-omni Platform. + + Inherits from vLLM's Platform and adds Omni-specific interfaces. + This gives OmniPlatform all vLLM Platform capabilities plus + Omni-specific methods. + """ + + _omni_enum: OmniPlatformEnum + + def is_npu(self) -> bool: + return self._omni_enum == OmniPlatformEnum.NPU + + @classmethod + def get_omni_ar_worker_cls(cls) -> str: + raise NotImplementedError + + @classmethod + def get_omni_generation_worker_cls(cls) -> str: + raise NotImplementedError + + @classmethod + def get_default_stage_config_path(cls) -> str: + raise NotImplementedError + + @classmethod + def get_diffusion_attn_backend_cls( + cls, + selected_backend: str | None, + head_size: int, + ) -> str: + """Get the diffusion attention backend class path for this platform. + + This method selects the appropriate attention backend for diffusion + models based on platform capabilities and user preferences. + + Args: + selected_backend: User-selected backend name (e.g., "FLASH_ATTN", + "TORCH_SDPA", "SAGE_ATTN"). If None, uses platform default. + head_size: Attention head size. + + Returns: + Fully qualified class path of the selected backend. + """ + raise NotImplementedError + + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + raise NotImplementedError + + @classmethod + def get_device_count(cls) -> int: + raise NotImplementedError + + @classmethod + def get_device_version(cls) -> str | None: + raise NotImplementedError + + @classmethod + def synchronize(cls) -> None: + raise NotImplementedError + + @classmethod + def get_free_memory(cls, device: torch.device | None = None) -> int: + raise NotImplementedError + + +class UnspecifiedOmniPlatform(OmniPlatform): + _omni_enum = OmniPlatformEnum.UNSPECIFIED + device_type = "" diff --git a/vllm_omni/platforms/npu/__init__.py b/vllm_omni/platforms/npu/__init__.py new file mode 100644 index 000000000..fb95f7b29 --- /dev/null +++ b/vllm_omni/platforms/npu/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.platforms.npu.platform import NPUOmniPlatform + +__all__ = ["NPUOmniPlatform"] diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py new file mode 100644 index 000000000..429297bbd --- /dev/null +++ b/vllm_omni/platforms/npu/platform.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.logger import init_logger +from vllm_ascend.platform import NPUPlatform + +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum + +logger = init_logger(__name__) + + +class NPUOmniPlatform(OmniPlatform, NPUPlatform): + """NPU/Ascend implementation of OmniPlatform. + + Inherits all NPU-specific implementations from vllm-ascend's NPUPlatform, + and adds Omni-specific interfaces from OmniPlatform. + """ + + _omni_enum = OmniPlatformEnum.NPU + dist_backend: str = "hccl" + + @classmethod + def get_omni_ar_worker_cls(cls) -> str: + return "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker" + + @classmethod + def get_omni_generation_worker_cls(cls) -> str: + return "vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker" + + @classmethod + def get_default_stage_config_path(cls) -> str: + return "vllm_omni/platforms/npu/stage_configs" + + @classmethod + def get_diffusion_attn_backend_cls( + cls, + selected_backend: str | None, + head_size: int, + ) -> str: + from importlib.util import find_spec + + if selected_backend is not None: + backend_upper = selected_backend.upper() + backend = DiffusionAttentionBackendEnum[backend_upper] + logger.info("Using diffusion attention backend '%s'", backend_upper) + return backend.get_path() + + # Try FLASH_ATTN if mindiesd is available, otherwise fall back to SDPA + if find_spec("mindiesd"): + logger.info("Defaulting to diffusion attention backend FLASH_ATTN") + return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + + logger.info("Falling back to diffusion attention backend SDPA") + return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + if local_rank is None: + return torch.device("npu") + return torch.device("npu", local_rank) + + @classmethod + def get_device_count(cls) -> int: + return torch.npu.device_count() + + @classmethod + def get_device_version(cls) -> str | None: + return None + + @classmethod + def synchronize(cls) -> None: + torch.npu.synchronize() + + @classmethod + def get_free_memory(cls, device: torch.device | None = None) -> int: + free, _ = torch.npu.mem_get_info(device) + return free + + @classmethod + def get_device_total_memory(cls, device_id: int = 0) -> int: + device_props = torch.npu.get_device_properties(device_id) + return device_props.total_memory diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml similarity index 89% rename from vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml rename to vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml index 49085276f..a02d84ed6 100644 --- a/vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen2_5_omni.yaml @@ -9,10 +9,10 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 - enforce_eager: true # Now we only support eager mode + enforce_eager: false trust_remote_code: true engine_output_type: latent enable_prefix_caching: false @@ -36,10 +36,10 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 - enforce_eager: true + enforce_eager: false trust_remote_code: true enable_prefix_caching: false engine_output_type: latent @@ -58,12 +58,12 @@ stage_args: stage_type: llm # Use llm stage type to launch OmniLLM runtime: process: true - devices: "0" # Example: use a different NPU than the previous stage; use "0" if single NPU + devices: "2" # Example: use a different NPU than the previous stage; use "0" if single NPU max_batch_size: 1 engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml similarity index 93% rename from vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml rename to vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml index eadfea570..fdf817239 100644 --- a/vllm_omni/model_executor/stage_configs/npu/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml @@ -12,7 +12,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: false @@ -41,7 +41,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.2 enforce_eager: false @@ -71,7 +71,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/model_executor/stage_configs/npu/qwen3_tts.yaml b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml similarity index 89% rename from vllm_omni/model_executor/stage_configs/npu/qwen3_tts.yaml rename to vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml index 86b5ca60e..d408dbab9 100644 --- a/vllm_omni/model_executor/stage_configs/npu/qwen3_tts.yaml +++ b/vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml @@ -7,7 +7,7 @@ stage_args: engine_args: model_stage: qwen3_tts model_arch: Qwen3TTSForConditionalGeneration - worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/worker/npu/__init__.py b/vllm_omni/platforms/npu/worker/__init__.py similarity index 100% rename from vllm_omni/worker/npu/__init__.py rename to vllm_omni/platforms/npu/worker/__init__.py diff --git a/vllm_omni/worker/npu/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py similarity index 99% rename from vllm_omni/worker/npu/npu_ar_model_runner.py rename to vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 1bb53b662..0ba57e038 100644 --- a/vllm_omni/worker/npu/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -37,7 +37,7 @@ from vllm_ascend.utils import ProfileExecuteDuration from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner +from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner class ExecuteModelState(NamedTuple): diff --git a/vllm_omni/worker/npu/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py similarity index 74% rename from vllm_omni/worker/npu/npu_ar_worker.py rename to vllm_omni/platforms/npu/worker/npu_ar_worker.py index c4895ee1b..1fa3e59ed 100644 --- a/vllm_omni/worker/npu/npu_ar_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_worker.py @@ -4,10 +4,11 @@ from vllm.v1.worker.workspace import init_workspace_manager from vllm_ascend.worker.worker import NPUWorker -from vllm_omni.worker.npu.npu_ar_model_runner import NPUARModelRunner +from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner +from vllm_omni.worker.mixins import OmniWorkerMixin -class NPUARWorker(NPUWorker): +class NPUARWorker(OmniWorkerMixin, NPUWorker): """NPU AR worker for thinker/talker stages in Omni model.""" def init_device(self): diff --git a/vllm_omni/worker/npu/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py similarity index 99% rename from vllm_omni/worker/npu/npu_generation_model_runner.py rename to vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index e1beef982..e8559bb46 100644 --- a/vllm_omni/worker/npu/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -27,8 +27,8 @@ from vllm_ascend.platform import NPUPlatform from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker.npu.npu_ar_model_runner import ExecuteModelState -from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner +from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState +from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner class NPUGenerationModelRunner(OmniNPUModelRunner): diff --git a/vllm_omni/worker/npu/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py similarity index 72% rename from vllm_omni/worker/npu/npu_generation_worker.py rename to vllm_omni/platforms/npu/worker/npu_generation_worker.py index 588d21fc1..edbaa9f64 100644 --- a/vllm_omni/worker/npu/npu_generation_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_worker.py @@ -4,10 +4,11 @@ from vllm.v1.worker.workspace import init_workspace_manager from vllm_ascend.worker.worker import NPUWorker -from vllm_omni.worker.npu.npu_generation_model_runner import NPUGenerationModelRunner +from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner +from vllm_omni.worker.mixins import OmniWorkerMixin -class NPUGenerationWorker(NPUWorker): +class NPUGenerationWorker(OmniWorkerMixin, NPUWorker): """NPU generation worker for code2wav stage in Omni model.""" def init_device(self): diff --git a/vllm_omni/worker/npu/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py similarity index 100% rename from vllm_omni/worker/npu/npu_model_runner.py rename to vllm_omni/platforms/npu/worker/npu_model_runner.py diff --git a/vllm_omni/platforms/rocm/__init__.py b/vllm_omni/platforms/rocm/__init__.py new file mode 100644 index 000000000..999dc5a08 --- /dev/null +++ b/vllm_omni/platforms/rocm/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.platforms.rocm.platform import RocmOmniPlatform + +__all__ = ["RocmOmniPlatform"] diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py new file mode 100644 index 000000000..f2c2d3a85 --- /dev/null +++ b/vllm_omni/platforms/rocm/platform.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.logger import init_logger +from vllm.platforms.rocm import RocmPlatform + +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum + +logger = init_logger(__name__) + + +class RocmOmniPlatform(OmniPlatform, RocmPlatform): + """ROCm/AMD GPU implementation of OmniPlatform. + + Inherits all ROCm-specific implementations from vLLM's RocmPlatform, + and adds Omni-specific interfaces from OmniPlatform. + """ + + _omni_enum = OmniPlatformEnum.ROCM + + @classmethod + def get_omni_ar_worker_cls(cls) -> str: + return "vllm_omni.worker.gpu_ar_worker.GPUARWorker" + + @classmethod + def get_omni_generation_worker_cls(cls) -> str: + return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker" + + @classmethod + def get_diffusion_attn_backend_cls( + cls, + selected_backend: str | None, + head_size: int, + ) -> str: + from vllm._aiter_ops import is_aiter_found_and_supported + + # Check if aiter is available for Flash Attention support + # aiter currently only is supported on gfx942 and gfx950 + # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + capability = major * 10 + minor + aiter_supported = is_aiter_found_and_supported() and 90 < capability < 100 + + if selected_backend is not None: + backend_upper = selected_backend.upper() + if backend_upper == "FLASH_ATTN" and not aiter_supported: + logger.warning( + "Flash Attention requires `aiter` library which is only supported " + "on gfx942 and gfx950. Falling back to TORCH_SDPA backend." + ) + logger.info("Defaulting to diffusion attention backend SDPA") + return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + backend = DiffusionAttentionBackendEnum[backend_upper] + logger.info("Using diffusion attention backend '%s'", backend_upper) + return backend.get_path() + + # Choose to enable Flash Attention by default on ROCm + # whenever possible as it is the fastest backend + if aiter_supported: + logger.info("Defaulting to diffusion attention backend FLASH_ATTN") + return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + + logger.info("Defaulting to diffusion attention backend SDPA") + return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + + @classmethod + def get_default_stage_config_path(cls) -> str: + return "vllm_omni/platforms/rocm/stage_configs" + + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + if local_rank is None: + return torch.device("cuda") + return torch.device("cuda", local_rank) + + @classmethod + def get_device_count(cls) -> int: + return torch.cuda.device_count() + + @classmethod + def get_device_version(cls) -> str | None: + if torch.version.hip is not None: + hip_version = torch.version.hip + return hip_version.split("-")[0] + return None + + @classmethod + def synchronize(cls) -> None: + torch.cuda.synchronize() + + @classmethod + def get_free_memory(cls, device: torch.device | None = None) -> int: + free, _ = torch.cuda.mem_get_info(device) + return free diff --git a/vllm_omni/model_executor/stage_configs/rocm/qwen2_5_omni.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml similarity index 94% rename from vllm_omni/model_executor/stage_configs/rocm/qwen2_5_omni.yaml rename to vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml index c646aa76a..7887cd2bb 100644 --- a/vllm_omni/model_executor/stage_configs/rocm/qwen2_5_omni.yaml +++ b/vllm_omni/platforms/rocm/stage_configs/qwen2_5_omni.yaml @@ -10,7 +10,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true # Now we only support eager mode @@ -38,7 +38,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.8 enforce_eager: true @@ -66,7 +66,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen2_5OmniForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler gpu_memory_utilization: 0.15 enforce_eager: true diff --git a/vllm_omni/model_executor/stage_configs/rocm/qwen3_omni_moe.yaml b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml similarity index 93% rename from vllm_omni/model_executor/stage_configs/rocm/qwen3_omni_moe.yaml rename to vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml index 73f65ecb5..93e05ea9c 100644 --- a/vllm_omni/model_executor/stage_configs/rocm/qwen3_omni_moe.yaml +++ b/vllm_omni/platforms/rocm/stage_configs/qwen3_omni_moe.yaml @@ -12,7 +12,7 @@ stage_args: engine_args: model_stage: thinker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 enforce_eager: true @@ -42,7 +42,7 @@ stage_args: engine_args: model_stage: talker model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 enforce_eager: true @@ -73,7 +73,7 @@ stage_args: engine_args: model_stage: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/platforms/xpu/__init__.py b/vllm_omni/platforms/xpu/__init__.py new file mode 100644 index 000000000..8ec0986c2 --- /dev/null +++ b/vllm_omni/platforms/xpu/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.platforms.xpu.platform import XPUOmniPlatform + +__all__ = ["XPUOmniPlatform"] diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py new file mode 100644 index 000000000..5c126c6d9 --- /dev/null +++ b/vllm_omni/platforms/xpu/platform.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.logger import init_logger +from vllm.platforms.xpu import XPUPlatform + +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum + +logger = init_logger(__name__) + + +class XPUOmniPlatform(OmniPlatform, XPUPlatform): + """XPU/Intel GPU implementation of OmniPlatform. + + Inherits all XPU-specific implementations from vLLM's XPUPlatform, + and adds Omni-specific interfaces from OmniPlatform. + """ + + _omni_enum = OmniPlatformEnum.XPU + + @classmethod + def get_omni_ar_worker_cls(cls) -> str: + return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker" + + @classmethod + def get_omni_generation_worker_cls(cls) -> str: + return "vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker" + + @classmethod + def get_diffusion_attn_backend_cls( + cls, + selected_backend: str | None, + head_size: int, + ) -> str: + if selected_backend is not None: + backend_upper = selected_backend.upper() + backend = DiffusionAttentionBackendEnum[backend_upper] + logger.info("Using diffusion attention backend '%s'", backend_upper) + return backend.get_path() + + logger.info("Defaulting to diffusion attention backend SDPA") + return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + + @classmethod + def get_default_stage_config_path(cls) -> str: + return "vllm_omni/platforms/xpu/stage_configs" + + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + if local_rank is None: + return torch.device("xpu") + return torch.device("xpu", local_rank) + + @classmethod + def get_device_count(cls) -> int: + return torch.xpu.device_count() + + @classmethod + def get_device_version(cls) -> str | None: + # XPU does not have a version string like CUDA + return None + + @classmethod + def synchronize(cls) -> None: + torch.xpu.synchronize() + + @classmethod + def get_free_memory(cls, device: torch.device | None = None) -> int: + if device is None: + device_id = 0 + else: + device_id = device.index if device.index is not None else 0 + props = torch.xpu.get_device_properties(device_id) + return props.total_memory diff --git a/vllm_omni/plugins/__init__.py b/vllm_omni/plugins/__init__.py new file mode 100644 index 000000000..41bf398ca --- /dev/null +++ b/vllm_omni/plugins/__init__.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +from collections.abc import Callable +from typing import Any + +import vllm.envs as envs + +logger = logging.getLogger(__name__) + +# Default plugins group will be loaded in all processes (process0, engine core +# process and worker processes). +OMNI_DEFAULT_PLUGINS_GROUP = "vllm_omni.general_plugins" +# Platform plugins group will be loaded in all processes when +# `vllm_omni.platforms.current_omni_platform` is called and the value not +# initialized. +OMNI_PLATFORM_PLUGINS_GROUP = "vllm_omni.platform_plugins" + +# Make sure one process only loads plugins once. +omni_plugins_loaded = False + + +def load_omni_plugins_by_group(group: str) -> dict[str, Callable[[], Any]]: + from importlib.metadata import entry_points + + allowed_plugins = envs.VLLM_PLUGINS + + discovered_plugins = entry_points(group=group) + if len(discovered_plugins) == 0: + logger.debug("No plugins for group %s found.", group) + return {} + + # Check if the only discovered plugin is the default one. + is_default_group = group == OMNI_DEFAULT_PLUGINS_GROUP + # Use INFO for non-default groups and DEBUG for the default group. + log_level = logger.debug if is_default_group else logger.info + + log_level("Available plugins for group %s:", group) + for plugin in discovered_plugins: + log_level("- %s -> %s", plugin.name, plugin.value) + + if allowed_plugins is None: + log_level("All plugins in this group will be loaded. Set `VLLM_PLUGINS` to control which plugins to load.") + + plugins: dict[str, Callable[[], Any]] = {} + for plugin in discovered_plugins: + if allowed_plugins is None or plugin.name in allowed_plugins: + if allowed_plugins is not None: + log_level("Loading plugin %s", plugin.name) + + try: + func = plugin.load() + plugins[plugin.name] = func + except Exception: + logger.exception("Failed to load plugin %s", plugin.name) + + return plugins + + +def load_omni_general_plugins() -> None: + """WARNING: plugins can be loaded for multiple times in different + processes. They should be designed in a way that they can be loaded + multiple times without causing issues. + """ + global omni_plugins_loaded + if omni_plugins_loaded: + return + omni_plugins_loaded = True + + plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP) + # General plugins: we only need to execute the loaded functions. + for func in plugins.values(): + func() + + +__all__ = [ + "load_omni_general_plugins", + "OMNI_DEFAULT_PLUGINS_GROUP", + "OMNI_PLATFORM_PLUGINS_GROUP", +] diff --git a/vllm_omni/utils/__init__.py b/vllm_omni/utils/__init__.py index 34b2545db..e69de29bb 100644 --- a/vllm_omni/utils/__init__.py +++ b/vllm_omni/utils/__init__.py @@ -1,13 +0,0 @@ -from vllm_omni.utils.platform_utils import ( - detect_device_type, - get_device_control_env_var, - is_npu, - is_rocm, -) - -__all__ = [ - "detect_device_type", - "get_device_control_env_var", - "is_npu", - "is_rocm", -] diff --git a/vllm_omni/utils/platform_utils.py b/vllm_omni/utils/platform_utils.py deleted file mode 100644 index fb47018e7..000000000 --- a/vllm_omni/utils/platform_utils.py +++ /dev/null @@ -1,58 +0,0 @@ -from __future__ import annotations - -import torch -from vllm.platforms import current_platform - - -def detect_device_type() -> str: - device_type = getattr(current_platform, "device_type", None) - if isinstance(device_type, str) and device_type: - return device_type.lower() - if torch.cuda.is_available(): - return "cuda" - if hasattr(torch, "npu") and torch.npu.is_available(): # type: ignore[attr-defined] - return "npu" - return "cpu" - - -def is_npu() -> bool: - return detect_device_type() == "npu" - - -def is_rocm() -> bool: - return current_platform.is_rocm() - - -def get_device_control_env_var() -> str: - """Return the environment variable name for device visibility control.""" - if hasattr(current_platform, "device_control_env_var"): - env_var = getattr(current_platform, "device_control_env_var", None) - if isinstance(env_var, str) and env_var: - return env_var - - device_type = detect_device_type() - if device_type == "npu": - return "ASCEND_RT_VISIBLE_DEVICES" - return "CUDA_VISIBLE_DEVICES" # fallback - - -def get_diffusion_worker_class() -> type: - """Get the appropriate diffusion WorkerProc class based on current device type. - - Returns: - The WorkerProc class for the detected device type (either NPUWorkerProc or WorkerProc). - - Raises: - ImportError: If the worker module for the detected device is not available. - """ - device_type = detect_device_type() - - if device_type == "npu": - from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorkerProc - - return NPUWorkerProc - else: - # Default to GPU worker for cuda and other devices - from vllm_omni.diffusion.worker.gpu_diffusion_worker import WorkerProc - - return WorkerProc diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index d2dafce68..a1d6b4665 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -13,11 +13,12 @@ from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner +from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) -class GPUARWorker(GPUWorker): +class GPUARWorker(OmniWorkerMixin, GPUWorker): """GPU worker for autoregressive omni model stages. Extends the base GPUWorker to initialize and manage autoregressive diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 19f8ab84b..0aa13ed76 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -13,11 +13,12 @@ from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner +from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) -class GPUGenerationWorker(GPUWorker): +class GPUGenerationWorker(OmniWorkerMixin, GPUWorker): """GPU Worker for Generation model (non-autoregressive waveform generation). Usage in stage config: diff --git a/vllm_omni/worker/mixins.py b/vllm_omni/worker/mixins.py new file mode 100644 index 000000000..5b25b8362 --- /dev/null +++ b/vllm_omni/worker/mixins.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from typing import Any + + +class OmniWorkerMixin: + """Mixin to ensure Omni plugins are loaded in worker processes.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + from vllm_omni.plugins import load_omni_general_plugins + + load_omni_general_plugins()