[Hardware] Support platforms and plugin system (#774)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
Canlin Guo authored and GitHub committed 2026-01-28 22:12:52 +08:00
1 parent b11d436e19
commit c9392dff01
103 files changed
+1770 -1218

No files matched your search

@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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:
@@ -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):
@@ -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:
@@ -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:
@@ -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
@@ -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
@@ -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.")
@@ -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
@@ -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
@@ -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
@@ -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
+8 -4
View File
@@ -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:
"""
+3 -3
View File
@@ -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