anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[MUSA][Feat] torch.accelerator support (#3101)

Browse Source 
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
This commit is contained in:
R0CKSTARandCanlin Guo authored and GitHub committed 2026-04-24 15:31:45 +00:00
1 parent 7b876324c2
commit d388583e0b
7 files changed
+9 -242
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/config/stage_config.py
+3 -18
View File
@@ -601,23 +601,6 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
return DeployConfig(**kwargs)
def _detect_platform() -> str | None:
"""Return "npu", "rocm", "xpu", or None (CUDA default)."""
try:
from vllm.platforms import current_platform
name = current_platform.device_name.lower()
if "npu" in name:
return "npu"
if "rocm" in name or "amd" in name:
return "rocm"
if "xpu" in name:
return "xpu"
except Exception as e:
logger.debug("Platform auto-detect failed, falling back to CUDA: %s", e)
return None
def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str | None]:
"""Return ``(overrides, devices)`` from a platform stage entry.
@@ -636,7 +619,9 @@ def _apply_platform_overrides(
) -> DeployConfig:
"""Merge platform-specific stage overrides into deploy config."""
if platform is None:
platform = _detect_platform()
from vllm_omni.platforms import current_omni_platform
platform = current_omni_platform.device_name.lower()
if platform is None or deploy.platforms is None:
return deploy
platform_section = deploy.platforms.get(platform)
vllm_omni/platforms/musa/platform.py
+2 -2
View File
@@ -25,11 +25,11 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
@classmethod
def get_omni_ar_worker_cls(cls) -> str:
return "vllm_omni.platforms.musa.worker.musa_ar_worker.MUSAARWorker"
return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"
@classmethod
def get_omni_generation_worker_cls(cls) -> str:
return "vllm_omni.platforms.musa.worker.musa_generation_worker.MUSAGenerationWorker"
return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
@classmethod
def get_default_stage_config_path(cls) -> str:
vllm_omni/platforms/musa/worker/__init__.py
-9
View File
@@ -1,9 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm_omni.platforms.musa.worker.musa_ar_worker import MUSAARWorker
from vllm_omni.platforms.musa.worker.musa_generation_worker import (
MUSAGenerationWorker,
)
__all__ = ["MUSAARWorker", "MUSAGenerationWorker"]
vllm_omni/platforms/musa/worker/musa_ar_worker.py
-103
View File
@@ -1,103 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MUSA AR (Autoregressive) worker for vLLM-Omni.
This worker handles autoregressive model stages (thinker/talker) on MUSA devices.
"""
import gc
import os
import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.mem_utils import MemorySnapshot, format_gib
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.utils import report_usage_stats
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm.v1.worker.utils import request_memory
from vllm.v1.worker.workspace import init_workspace_manager
from vllm_omni.worker.base import OmniGPUWorkerBase
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
from vllm_omni.worker.mixins import OmniWorkerMixin
logger = init_logger(__name__)
class MUSAARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
"""MUSA AR worker for thinker/talker stages in Omni model."""
def init_device(self):
"""Initialize the MUSA device for this worker."""
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
if (
parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
and parallel_config.data_parallel_backend != "ray"
and parallel_config.nnodes_within_dp == 1
):
# Use local DP rank if available, otherwise use global DP rank.
dp_local_rank = self.parallel_config.data_parallel_rank_local
if dp_local_rank is None:
dp_local_rank = self.parallel_config.data_parallel_index
tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
self.local_rank += dp_local_rank * tp_pp_world_size
assert self.local_rank < torch.musa.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.musa.device_count()
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
f"({visible_device_count})."
)
self.device = torch.device(f"musa:{self.local_rank}")
torch.musa.set_device(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)
# Initialize the distributed environment BEFORE taking memory snapshot
# This ensures NCCL buffers are allocated before we measure available memory
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
# Set random seed.
set_random_seed(self.model_config.seed)
# Now take memory snapshot after distributed environment is initialized
gc.collect()
torch.musa.empty_cache()
# Take current memory snapshot
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
# Initialize workspace manager
num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
init_workspace_manager(self.device, num_ubatches)
if self.use_v2_model_runner:
# OMNI: v2 model runner does not yet include omni hooks.
logger.warning("OMNI MUSAARWorker forces v1 model runner for omni hooks.")
self.use_v2_model_runner = False
# Construct the model runner
self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
report_usage_stats(self.vllm_config)
vllm_omni/platforms/musa/worker/musa_generation_worker.py
-106
View File
@@ -1,106 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""MUSA generation worker for vLLM-Omni.
This worker handles non-autoregressive generation stages (e.g., code2wav waveform
generation) on MUSA devices.
"""
import gc
import os
import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.tracing import instrument
from vllm.utils.mem_utils import MemorySnapshot, format_gib
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.utils import report_usage_stats
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm.v1.worker.utils import request_memory
from vllm.v1.worker.workspace import init_workspace_manager
from vllm_omni.worker.base import OmniGPUWorkerBase
from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
from vllm_omni.worker.mixins import OmniWorkerMixin
logger = init_logger(__name__)
class MUSAGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
"""MUSA generation worker for non-AR waveform generation stage."""
@instrument(span_name="Init device")
def init_device(self):
"""Initialize the MUSA device for this worker."""
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
if (
parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
and parallel_config.data_parallel_backend != "ray"
and parallel_config.nnodes_within_dp == 1
):
# Use local DP rank if available, otherwise use global DP rank.
dp_local_rank = self.parallel_config.data_parallel_rank_local
if dp_local_rank is None:
dp_local_rank = self.parallel_config.data_parallel_index
tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
self.local_rank += dp_local_rank * tp_pp_world_size
assert self.local_rank < torch.musa.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.musa.device_count()
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
f"({visible_device_count})."
)
self.device = torch.device(f"musa:{self.local_rank}")
torch.musa.set_device(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)
# Initialize the distributed environment BEFORE taking memory snapshot
# This ensures NCCL buffers are allocated before we measure available memory
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
# Set random seed.
set_random_seed(self.model_config.seed)
# Now take memory snapshot after distributed environment is initialized
gc.collect()
torch.musa.empty_cache()
# Take current memory snapshot
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
# Initialize workspace manager
num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
init_workspace_manager(self.device, num_ubatches)
if self.use_v2_model_runner:
# OMNI: v2 model runner does not yet include omni hooks.
logger.warning("OMNI MUSAGenerationWorker forces v1 model runner for omni hooks.")
self.use_v2_model_runner = False
# Construct the model runner
self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
report_usage_stats(self.vllm_config)
vllm_omni/worker/gpu_ar_worker.py
+2 -2
View File
@@ -29,7 +29,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
@instrument(span_name="Init device")
def init_device(self):
if self.device_config.device_type == "cuda":
if self.device_config.device_type in ("cuda", "musa"):
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
@@ -52,7 +52,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
assert self.local_rank < torch.accelerator.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
visible_device_count = torch.accelerator.device_count()
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "
vllm_omni/worker/gpu_generation_worker.py
+2 -2
View File
@@ -28,7 +28,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
@instrument(span_name="Init device")
def init_device(self):
if self.device_config.device_type == "cuda":
if self.device_config.device_type in ("cuda", "musa"):
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
parallel_config = self.parallel_config
@@ -51,7 +51,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
assert self.local_rank < torch.accelerator.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of bounds. "
)
visible_device_count = torch.accelerator.device_count() if torch.cuda.is_available() else 0
visible_device_count = torch.accelerator.device_count()
assert self.parallel_config.local_world_size <= visible_device_count, (
f"local_world_size ({self.parallel_config.local_world_size}) must "
f"be less than or equal to the number of visible devices "