wehub-resource-sync
1011 lines
36 KiB
Python
1011 lines
36 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Code inside this file can safely assume cuda platform, e.g. importing
|
|
pynvml. However, it should not initialize cuda context.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import os
|
|
import platform
|
|
from collections.abc import Callable
|
|
from datetime import timedelta
|
|
from functools import cache, lru_cache, wraps
|
|
from typing import TYPE_CHECKING, NamedTuple, TypeVar
|
|
|
|
import torch
|
|
from torch.distributed import PrefixStore, ProcessGroup
|
|
from torch.distributed.distributed_c10d import is_nccl_available
|
|
from typing_extensions import ParamSpec
|
|
|
|
# import custom ops, trigger op registration
|
|
import vllm._C_stable_libtorch # noqa
|
|
|
|
with contextlib.suppress(ImportError):
|
|
import vllm._qutlass_C # noqa
|
|
import vllm.envs as envs
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.import_utils import import_pynvml
|
|
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
|
|
|
from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.cache import CacheDType
|
|
from vllm.config.kernel import IrOpPriorityConfig
|
|
from vllm.v1.attention.backend import AttentionBackend
|
|
from vllm.v1.attention.selector import AttentionSelectorConfig
|
|
else:
|
|
VllmConfig = None
|
|
CacheDType = None
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_P = ParamSpec("_P")
|
|
_R = TypeVar("_R")
|
|
|
|
pynvml = import_pynvml()
|
|
|
|
# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
|
|
# see https://github.com/huggingface/diffusers/issues/9704 for details
|
|
torch.backends.cuda.enable_cudnn_sdp(False)
|
|
|
|
|
|
@lru_cache(maxsize=8)
|
|
def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:
|
|
"""Get number of CUDA devices, caching based on the value of CUDA_VISIBLE_DEVICES
|
|
at the time of call.
|
|
|
|
This should be used instead of torch.accelerator.device_count() unless
|
|
CUDA_VISIBLE_DEVICES has already been set to the desired value.
|
|
|
|
# This can be removed and simply replaced with torch.cuda.get_device_count
|
|
# after https://github.com/pytorch/pytorch/pull/122815 is released."""
|
|
# Note: cuda_visible_devices is not used, but we keep it as an argument for
|
|
# LRU Cache purposes.
|
|
|
|
# Code below is based on
|
|
# https://github.com/pytorch/pytorch/blob/
|
|
# c1cd946818442aca8c7f812b16d187ce1586c3bc/
|
|
# torch/cuda/__init__.py#L831C1-L831C17
|
|
import torch.cuda
|
|
|
|
if not torch.cuda._is_compiled():
|
|
return 0
|
|
raw_count = torch.cuda._device_count_nvml()
|
|
r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
|
|
return r
|
|
|
|
|
|
@cache
|
|
def _get_backend_priorities(
|
|
use_mla: bool,
|
|
device_capability: DeviceCapability,
|
|
num_heads: int | None = None,
|
|
kv_cache_dtype: CacheDType | None = None,
|
|
) -> list[AttentionBackendEnum]:
|
|
"""Get backend priorities with lazy import to avoid circular dependency."""
|
|
from vllm.utils.torch_utils import is_quantized_kv_cache
|
|
|
|
if use_mla:
|
|
if device_capability.major == 10:
|
|
# Sparse MLA backend priorities
|
|
# See https://github.com/vllm-project/vllm/issues/35807 for
|
|
# benchmark results
|
|
if kv_cache_dtype is not None and is_quantized_kv_cache(kv_cache_dtype):
|
|
# Prefer FlashInfer for fp8 kv cache
|
|
sparse_backends = [
|
|
AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
|
|
AttentionBackendEnum.FLASHMLA_SPARSE,
|
|
]
|
|
else:
|
|
# BF16 KV Cache
|
|
# Prefer FlashInfer at low head counts (FlashMLA uses padding)
|
|
if num_heads is not None and num_heads <= 16:
|
|
sparse_backends = [
|
|
AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
|
|
AttentionBackendEnum.FLASHMLA_SPARSE,
|
|
]
|
|
else:
|
|
sparse_backends = [
|
|
AttentionBackendEnum.FLASHMLA_SPARSE,
|
|
AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
|
|
]
|
|
|
|
return [
|
|
AttentionBackendEnum.FLASHINFER_MLA,
|
|
# R1 dims + FP8 KV only; rejected by supports_combination
|
|
# otherwise. Behind FLASHINFER_MLA: wins past bs≈8, regresses
|
|
# at bs≤2.
|
|
AttentionBackendEnum.TOKENSPEED_MLA,
|
|
AttentionBackendEnum.CUTLASS_MLA,
|
|
AttentionBackendEnum.FLASH_ATTN_MLA,
|
|
AttentionBackendEnum.FLASHMLA,
|
|
AttentionBackendEnum.TRITON_MLA,
|
|
*sparse_backends,
|
|
]
|
|
elif device_capability.major == 12:
|
|
return [
|
|
AttentionBackendEnum.TRITON_MLA,
|
|
AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
|
|
]
|
|
else:
|
|
return [
|
|
AttentionBackendEnum.FLASH_ATTN_MLA,
|
|
AttentionBackendEnum.FLASHMLA,
|
|
AttentionBackendEnum.FLASHINFER_MLA,
|
|
AttentionBackendEnum.TRITON_MLA,
|
|
AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
|
|
AttentionBackendEnum.FLASHMLA_SPARSE,
|
|
]
|
|
else:
|
|
if device_capability.major == 10:
|
|
return [
|
|
AttentionBackendEnum.FLASHINFER,
|
|
AttentionBackendEnum.FLASH_ATTN,
|
|
AttentionBackendEnum.TRITON_ATTN,
|
|
AttentionBackendEnum.FLEX_ATTENTION,
|
|
AttentionBackendEnum.TURBOQUANT,
|
|
]
|
|
else:
|
|
return [
|
|
AttentionBackendEnum.FLASH_ATTN,
|
|
AttentionBackendEnum.FLASHINFER,
|
|
AttentionBackendEnum.TRITON_ATTN,
|
|
AttentionBackendEnum.FLEX_ATTENTION,
|
|
AttentionBackendEnum.TURBOQUANT,
|
|
]
|
|
|
|
|
|
def _backend_cls_path(backend_cls: type[AttentionBackend]) -> str:
|
|
module, qualname = backend_cls.full_cls_name()
|
|
return f"{module}.{qualname}"
|
|
|
|
|
|
def _get_attn_backend_class(backend: AttentionBackendEnum) -> type[AttentionBackend]:
|
|
return backend.get_class()
|
|
|
|
|
|
class _BackendCandidate(NamedTuple):
|
|
backend_class: type[AttentionBackend]
|
|
backend: AttentionBackendEnum
|
|
priority: int
|
|
|
|
|
|
def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
|
|
@wraps(fn)
|
|
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
|
pynvml.nvmlInit()
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
finally:
|
|
pynvml.nvmlShutdown()
|
|
|
|
return wrapper
|
|
|
|
|
|
@cache
|
|
def _get_wsl_kernel_version() -> tuple[int, ...] | None:
|
|
"""Return the WSL2 kernel version as a tuple, or None on parse failure.
|
|
|
|
platform.uname().release on WSL2 looks like
|
|
"5.15.167.4-microsoft-standard-WSL2"; we take the numeric prefix.
|
|
"""
|
|
try:
|
|
release = platform.uname().release
|
|
parts = release.split("-")[0].split(".")
|
|
return tuple(int(x) for x in parts[:3])
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
class CudaPlatformBase(Platform):
|
|
_enum = PlatformEnum.CUDA
|
|
device_name: str = "cuda"
|
|
device_type: str = "cuda"
|
|
dispatch_key: str = "CUDA"
|
|
ray_device_key: str = "GPU"
|
|
dist_backend: str = "nccl"
|
|
device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
|
|
ray_noset_device_env_vars: list[str] = [
|
|
"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
|
|
]
|
|
|
|
@classmethod
|
|
def import_kernels(cls) -> None:
|
|
"""Import CUDA kernel extensions (_C_stable_libtorch, optional _qutlass_C)."""
|
|
try:
|
|
import vllm._C_stable_libtorch # noqa: F401
|
|
except ImportError as e:
|
|
logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e)
|
|
with contextlib.suppress(ImportError):
|
|
import vllm._moe_C_stable_libtorch # noqa: F401
|
|
with contextlib.suppress(ImportError):
|
|
import vllm._qutlass_C # noqa: F401
|
|
|
|
@property
|
|
def supported_dtypes(self) -> list[torch.dtype]:
|
|
if self.has_device_capability(80):
|
|
# Ampere and Hopper or later NVIDIA GPUs.
|
|
return [torch.bfloat16, torch.float16, torch.float32]
|
|
if self.has_device_capability(60):
|
|
# Pascal, Volta and Turing NVIDIA GPUs, BF16 is not supported
|
|
return [torch.float16, torch.float32]
|
|
# Kepler and Maxwell NVIDIA GPUs, only FP32 is supported,
|
|
# though vLLM doesn't support these GPUs.
|
|
return [torch.float32]
|
|
|
|
@classmethod
|
|
def set_device(cls, device: torch.device) -> None:
|
|
"""
|
|
Set the device for the current platform.
|
|
"""
|
|
torch.cuda.set_device(device)
|
|
# With this trick we can force the device to be set eagerly
|
|
# see https://github.com/pytorch/pytorch/issues/155668
|
|
# for why and when it is needed
|
|
_ = torch.zeros(1, device=device)
|
|
|
|
@classmethod
|
|
def manual_seed_all(cls, seed: int) -> None:
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
@classmethod
|
|
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def get_cuda_runtime_major(cls) -> int:
|
|
"""Major ``torch.version.cuda`` version, or ``0`` if undetermined."""
|
|
major = (torch.version.cuda or "0").split(".", 1)[0]
|
|
return int(major) if major.isdigit() else 0
|
|
|
|
@classmethod
|
|
def get_device_name(cls, device_id: int = 0) -> str:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def get_device_total_memory(cls, device_id: int = 0) -> int:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def is_fully_connected(cls, device_ids: list[int]) -> bool:
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def log_warnings(cls):
|
|
pass
|
|
|
|
@classmethod
|
|
def is_pin_memory_available(cls) -> bool:
|
|
if in_wsl():
|
|
# WSL1 has no CUDA support, so being on the CUDA platform under
|
|
# WSL implies WSL2. Gate on kernel >= 4.19.121, the first WSL2
|
|
# kernel with limited pinned memory support for CUDA.
|
|
version = _get_wsl_kernel_version()
|
|
if version is None or version < (4, 19, 121):
|
|
logger.warning_once(
|
|
"Using 'pin_memory=False' as WSL is detected and the "
|
|
"WSL2 kernel version is below 4.19.121. This may slow "
|
|
"down performance. Please run `wsl --update`."
|
|
)
|
|
return False
|
|
# On compatible WSL2 kernels, pinned memory is supported but
|
|
# disabled by default. Enable it via VLLM_WSL2_ENABLE_PIN_MEMORY=1.
|
|
import vllm.envs as envs
|
|
|
|
return envs.VLLM_WSL2_ENABLE_PIN_MEMORY
|
|
return True
|
|
|
|
@classmethod
|
|
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
|
|
parallel_config = vllm_config.parallel_config
|
|
model_config = vllm_config.model_config
|
|
|
|
if parallel_config.worker_cls == "auto":
|
|
parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
|
|
|
|
scheduler_config = vllm_config.scheduler_config
|
|
# Note: model_config may be None during testing
|
|
if (
|
|
model_config is not None
|
|
and model_config.is_mm_prefix_lm
|
|
and scheduler_config.is_multimodal_model
|
|
and not scheduler_config.disable_chunked_mm_input
|
|
):
|
|
logger.warning_once(
|
|
"Forcing --disable_chunked_mm_input for models "
|
|
"with multimodal-bidirectional attention."
|
|
)
|
|
scheduler_config.disable_chunked_mm_input = True
|
|
|
|
if (
|
|
in_wsl()
|
|
and vllm_config.offload_config.uva.cpu_offload_gb > 0
|
|
and bool(vllm_config.compilation_config.cudagraph_mode)
|
|
):
|
|
logger.warning_once(
|
|
"--cpu-offload-gb is enabled with CUDA graphs on WSL2. "
|
|
"This combination requires pinned (page-locked) memory "
|
|
"allocations. WARNING: Windows (WDDM) enforces a hard "
|
|
"system-wide cap of roughly 50%% of physical RAM on pinned "
|
|
"memory shared across ALL processes by default (limit can "
|
|
"changed via %%USERPROFILE%%\\.wslconfig). "
|
|
"Excessive use of page-locked memory can prevent Windows "
|
|
"from reclaiming memory under load, which can cause the "
|
|
"entire host OS to become unresponsive and may require a "
|
|
"hard reboot to recover. Proceed at your own risk. "
|
|
"To raise the WSL2 VM memory ceiling, increase the `memory` "
|
|
"setting in %%USERPROFILE%%\\.wslconfig and run "
|
|
"`wsl --shutdown`."
|
|
)
|
|
|
|
@classmethod
|
|
def get_current_memory_usage(
|
|
cls, device: torch.types.Device | None = None
|
|
) -> float:
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
return torch.cuda.max_memory_allocated(device)
|
|
|
|
@classmethod
|
|
def get_valid_backends(
|
|
cls,
|
|
device_capability: DeviceCapability,
|
|
attn_selector_config: AttentionSelectorConfig,
|
|
num_heads: int | None = None,
|
|
) -> tuple[
|
|
list[_BackendCandidate],
|
|
dict[AttentionBackendEnum, tuple[int, list[str]]],
|
|
]:
|
|
valid_backends_priorities = []
|
|
invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
|
|
|
|
backend_priorities = _get_backend_priorities(
|
|
attn_selector_config.use_mla,
|
|
device_capability,
|
|
num_heads,
|
|
attn_selector_config.kv_cache_dtype,
|
|
)
|
|
for priority, backend in enumerate(backend_priorities):
|
|
try:
|
|
backend_class = _get_attn_backend_class(backend)
|
|
invalid_reasons_i = backend_class.validate_configuration(
|
|
device_capability=device_capability,
|
|
**attn_selector_config._asdict(),
|
|
)
|
|
except ImportError:
|
|
invalid_reasons_i = ["ImportError"]
|
|
if invalid_reasons_i:
|
|
invalid_reasons[backend] = (priority, invalid_reasons_i)
|
|
else:
|
|
valid_backends_priorities.append(
|
|
_BackendCandidate(backend_class, backend, priority)
|
|
)
|
|
|
|
return valid_backends_priorities, invalid_reasons
|
|
|
|
@classmethod
|
|
def get_attn_backend_cls(
|
|
cls,
|
|
selected_backend: AttentionBackendEnum | None,
|
|
attn_selector_config: AttentionSelectorConfig,
|
|
num_heads: int | None = None,
|
|
) -> str:
|
|
device_capability = cls.get_device_capability()
|
|
assert device_capability is not None
|
|
|
|
# First try checking just the selected backend, if there is one.
|
|
if selected_backend is not None:
|
|
try:
|
|
backend_class = _get_attn_backend_class(selected_backend)
|
|
invalid_reasons = backend_class.validate_configuration(
|
|
device_capability=device_capability,
|
|
**attn_selector_config._asdict(),
|
|
)
|
|
except ImportError:
|
|
invalid_reasons = ["ImportError"]
|
|
if invalid_reasons:
|
|
raise ValueError(
|
|
f"Selected backend {selected_backend} is not valid for "
|
|
f"this configuration. Reason: {invalid_reasons}"
|
|
)
|
|
else:
|
|
logger.info("Using %s backend.", selected_backend)
|
|
return _backend_cls_path(backend_class)
|
|
|
|
# No selected backend or the selected backend is invalid,
|
|
# so we try finding a valid backend.
|
|
valid_backends_priorities, all_invalid_reasons = cls.get_valid_backends(
|
|
device_capability=device_capability,
|
|
attn_selector_config=attn_selector_config,
|
|
num_heads=num_heads,
|
|
)
|
|
reasons_str = (
|
|
"{"
|
|
+ ", ".join(
|
|
f"{backend.name}: [{', '.join(reasons)}]"
|
|
for backend, (_, reasons) in all_invalid_reasons.items()
|
|
)
|
|
+ "}"
|
|
)
|
|
config_str = attn_selector_config.__repr__()
|
|
logger.debug_once(
|
|
f"Some attention backends are not valid for {cls.device_name} with "
|
|
f"{config_str}. Reasons: {reasons_str}."
|
|
)
|
|
if len(valid_backends_priorities) == 0:
|
|
raise ValueError(
|
|
f"No valid attention backend found for {cls.device_name} "
|
|
f"with {config_str}. Reasons: {reasons_str}."
|
|
)
|
|
|
|
# We have found some valid backends. Select the one with the
|
|
# highest priority.
|
|
selected_candidate = min(
|
|
valid_backends_priorities,
|
|
key=lambda candidate: candidate.priority,
|
|
)
|
|
selected_backend_class = selected_candidate.backend_class
|
|
selected_backend = selected_candidate.backend
|
|
selected_priority = selected_candidate.priority
|
|
|
|
# If the user specified --block-size (but not --attention-backend),
|
|
# check whether that constraint precluded any higher-priority backends.
|
|
if attn_selector_config.block_size is not None:
|
|
excluded = [
|
|
backend
|
|
for backend, (priority, reasons) in all_invalid_reasons.items()
|
|
if priority < selected_priority
|
|
and reasons == ["block_size not supported"]
|
|
]
|
|
if excluded:
|
|
names = ", ".join(b.name for b in excluded)
|
|
logger.warning(
|
|
"--block-size %d precluded higher-priority backend(s) "
|
|
"%s. Using %s instead, which may result in reduced "
|
|
"performance. Consider removing --block-size to "
|
|
"auto-select the optimal block size.",
|
|
attn_selector_config.block_size,
|
|
names,
|
|
selected_backend.name,
|
|
)
|
|
|
|
logger.info_once(
|
|
"Using %s attention backend out of potential backends: %s.",
|
|
selected_backend.name,
|
|
"["
|
|
+ ", ".join(
|
|
f"'{candidate.backend.name}'" for candidate in valid_backends_priorities
|
|
)
|
|
+ "]",
|
|
)
|
|
|
|
return _backend_cls_path(selected_backend_class)
|
|
|
|
@classmethod
|
|
def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]:
|
|
if cls.has_device_capability(80):
|
|
return [
|
|
AttentionBackendEnum.FLASH_ATTN,
|
|
AttentionBackendEnum.TRITON_ATTN,
|
|
AttentionBackendEnum.TORCH_SDPA,
|
|
AttentionBackendEnum.FLASHINFER,
|
|
]
|
|
else:
|
|
return [
|
|
AttentionBackendEnum.FLASH_ATTN,
|
|
AttentionBackendEnum.TORCH_SDPA,
|
|
AttentionBackendEnum.TRITON_ATTN,
|
|
AttentionBackendEnum.FLASHINFER,
|
|
]
|
|
|
|
@classmethod
|
|
def get_vit_attn_backend(
|
|
cls,
|
|
head_size: int,
|
|
dtype: torch.dtype,
|
|
backend: AttentionBackendEnum | None = None,
|
|
) -> AttentionBackendEnum:
|
|
if backend is not None:
|
|
assert backend in cls.get_supported_vit_attn_backends(), (
|
|
f"Backend {backend} is not supported for vit attention. "
|
|
f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
|
|
)
|
|
logger.info_once(f"Using backend {backend} for vit attention")
|
|
return backend
|
|
|
|
cc = cls.get_device_capability()
|
|
for vit_attn_backend in cls.get_supported_vit_attn_backends():
|
|
if vit_attn_backend == AttentionBackendEnum.TORCH_SDPA:
|
|
return vit_attn_backend
|
|
try:
|
|
backend_class = vit_attn_backend.get_class()
|
|
is_backend_supported = backend_class.supports_head_size(
|
|
head_size
|
|
) and backend_class.supports_dtype(dtype)
|
|
if cc is not None:
|
|
is_backend_supported = (
|
|
is_backend_supported
|
|
and backend_class.supports_compute_capability(cc)
|
|
)
|
|
if is_backend_supported:
|
|
logger.info_once(
|
|
f"Using backend {vit_attn_backend} for vit attention",
|
|
)
|
|
return vit_attn_backend
|
|
except ImportError:
|
|
pass
|
|
|
|
return AttentionBackendEnum.TORCH_SDPA
|
|
|
|
@classmethod
|
|
def get_punica_wrapper(cls) -> str:
|
|
return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
|
|
|
|
@classmethod
|
|
def get_device_communicator_cls(cls) -> str:
|
|
return (
|
|
"vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa
|
|
)
|
|
|
|
@classmethod
|
|
def supports_fp8(cls) -> bool:
|
|
return cls.has_device_capability(89)
|
|
|
|
@classmethod
|
|
def use_custom_allreduce(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def opaque_attention_op(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def get_static_graph_wrapper_cls(cls) -> str:
|
|
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
|
|
|
|
@classmethod
|
|
def stateless_init_device_torch_dist_pg(
|
|
cls,
|
|
backend: str,
|
|
prefix_store: PrefixStore,
|
|
group_rank: int,
|
|
group_size: int,
|
|
timeout: timedelta,
|
|
) -> ProcessGroup:
|
|
assert is_nccl_available()
|
|
pg: ProcessGroup = ProcessGroup(
|
|
prefix_store,
|
|
group_rank,
|
|
group_size,
|
|
)
|
|
from torch.distributed.distributed_c10d import ProcessGroupNCCL
|
|
|
|
backend_options = ProcessGroupNCCL.Options()
|
|
backend_options._timeout = timeout
|
|
|
|
backend_class = ProcessGroupNCCL(
|
|
prefix_store, group_rank, group_size, backend_options
|
|
)
|
|
backend_type = ProcessGroup.BackendType.NCCL
|
|
device = torch.device("cuda")
|
|
pg._set_default_backend(backend_type)
|
|
backend_class._set_sequence_number_for_group()
|
|
|
|
pg._register_backend(device, backend_type, backend_class)
|
|
return pg
|
|
|
|
@classmethod
|
|
def device_count(cls) -> int:
|
|
return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)
|
|
|
|
@classmethod
|
|
def check_if_supports_dtype(cls, dtype: torch.dtype):
|
|
if dtype == torch.bfloat16: # noqa: SIM102
|
|
if not cls.has_device_capability(80):
|
|
capability = cls.get_device_capability()
|
|
gpu_name = cls.get_device_name()
|
|
|
|
if capability is None:
|
|
compute_str = "does not have a compute capability"
|
|
else:
|
|
version_str = capability.as_version_str()
|
|
compute_str = f"has compute capability {version_str}"
|
|
|
|
raise ValueError(
|
|
"Bfloat16 is only supported on GPUs "
|
|
"with compute capability of at least 8.0. "
|
|
f"Your {gpu_name} GPU {compute_str}. "
|
|
"You can use float16 instead by explicitly setting the "
|
|
"`dtype` flag in CLI, for example: --dtype=half."
|
|
)
|
|
|
|
@classmethod
|
|
def insert_blocks_to_device(
|
|
cls,
|
|
src_cache: torch.Tensor,
|
|
dst_cache: torch.Tensor,
|
|
src_block_indices: torch.Tensor,
|
|
dst_block_indices: torch.Tensor,
|
|
) -> None:
|
|
"""Copy blocks from src_cache to dst_cache on GPU."""
|
|
_src_cache = src_cache[src_block_indices]
|
|
dst_cache[dst_block_indices] = _src_cache.to(dst_cache.device)
|
|
|
|
@classmethod
|
|
def swap_out_blocks_to_host(
|
|
cls,
|
|
src_cache: torch.Tensor,
|
|
dst_cache: torch.Tensor,
|
|
src_block_indices: torch.Tensor,
|
|
dst_block_indices: torch.Tensor,
|
|
) -> None:
|
|
"""Copy blocks from GPU to host (CPU)."""
|
|
_src_cache = src_cache[src_block_indices]
|
|
dst_cache[dst_block_indices] = _src_cache.cpu()
|
|
|
|
@classmethod
|
|
def support_hybrid_kv_cache(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def support_static_graph_mode(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def support_deep_gemm(cls) -> bool:
|
|
"""Currently, only Hopper and Blackwell GPUs are supported."""
|
|
return (
|
|
cls.is_device_capability(90)
|
|
or cls.is_device_capability_family(100)
|
|
or cls.is_device_capability_family(120)
|
|
)
|
|
|
|
@classmethod
|
|
def is_integrated_gpu(cls, device_id: int = 0) -> bool:
|
|
return bool(torch.cuda.get_device_properties(device_id).is_integrated)
|
|
|
|
@classmethod
|
|
def num_compute_units(cls, device_id: int = 0) -> int:
|
|
return torch.cuda.get_device_properties(device_id).multi_processor_count
|
|
|
|
@classmethod
|
|
def use_custom_op_collectives(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
|
|
from vllm.config.compilation import CompilationMode
|
|
from vllm.config.kernel import IrOpPriorityConfig
|
|
|
|
# Native used by default when compiling,
|
|
# use vllm_c kernels where available when no codegen
|
|
cc = vllm_config.compilation_config
|
|
using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
|
|
default = ["native"] if using_inductor else ["vllm_c", "native"]
|
|
|
|
# Use oink if enabled for rms_norm
|
|
# TODO(Laurawly/luka): remove this env var,
|
|
# users can just use IR op priority directly
|
|
rms_norm = default
|
|
if envs.VLLM_USE_OINK_OPS:
|
|
rms_norm = ["oink"] + default
|
|
|
|
return IrOpPriorityConfig.with_default(
|
|
default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
|
|
)
|
|
|
|
@classmethod
|
|
def is_arch_support_pdl(cls) -> bool:
|
|
try:
|
|
device = torch.cuda.current_device()
|
|
major, _ = torch.cuda.get_device_capability(device)
|
|
except Exception:
|
|
return False
|
|
return major >= 9
|
|
|
|
|
|
# NVML utils
|
|
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
|
|
# all the related functions work on real physical device ids.
|
|
# the major benefit of using NVML is that it will not initialize CUDA
|
|
class NvmlCudaPlatform(CudaPlatformBase):
|
|
@classmethod
|
|
@with_nvml_context
|
|
def device_control_id_to_physical_device_id(cls, device_id: str) -> int:
|
|
try:
|
|
return int(device_id)
|
|
except ValueError:
|
|
handle = pynvml.nvmlDeviceGetHandleByUUID(device_id)
|
|
return pynvml.nvmlDeviceGetIndex(handle)
|
|
|
|
@classmethod
|
|
@cache
|
|
@with_nvml_context
|
|
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
|
|
try:
|
|
physical_device_id = cls.visible_device_id_to_physical_device_id(device_id)
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
|
|
major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
|
|
return DeviceCapability(major=major, minor=minor)
|
|
except RuntimeError:
|
|
return None
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def has_device_capability(
|
|
cls,
|
|
capability: tuple[int, int] | int,
|
|
device_id: int = 0,
|
|
) -> bool:
|
|
try:
|
|
return super().has_device_capability(capability, device_id)
|
|
except RuntimeError:
|
|
return False
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def get_device_name(cls, device_id: int = 0) -> str:
|
|
physical_device_id = cls.device_id_to_physical_device_id(device_id)
|
|
return cls._get_physical_device_name(physical_device_id)
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def get_device_uuid(cls, device_id: int = 0) -> str:
|
|
physical_device_id = cls.device_id_to_physical_device_id(device_id)
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
|
|
return pynvml.nvmlDeviceGetUUID(handle)
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def get_device_total_memory(cls, device_id: int = 0) -> int:
|
|
physical_device_id = cls.device_id_to_physical_device_id(device_id)
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
|
|
return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
|
|
"""
|
|
query if the set of gpus are fully connected by nvlink (1 hop)
|
|
"""
|
|
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
|
|
for i, handle in enumerate(handles):
|
|
for j, peer_handle in enumerate(handles):
|
|
if i < j:
|
|
try:
|
|
p2p_status = pynvml.nvmlDeviceGetP2PStatus(
|
|
handle,
|
|
peer_handle,
|
|
pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
|
|
)
|
|
if p2p_status != pynvml.NVML_P2P_STATUS_OK:
|
|
return False
|
|
except pynvml.NVMLError:
|
|
logger.exception(
|
|
"NVLink detection failed. This is normal if"
|
|
" your machine has no NVLink equipped."
|
|
)
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def _get_physical_device_name(cls, device_id: int = 0) -> str:
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
|
|
return pynvml.nvmlDeviceGetName(handle)
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def get_device_numa_node(cls, device_id: int = 0) -> int | None:
|
|
"""Get the NUMA node ID for a GPU device."""
|
|
physical_device_id = cls.device_id_to_physical_device_id(device_id)
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
|
|
|
|
try:
|
|
numa_node = pynvml.nvmlDeviceGetNumaNodeId(handle)
|
|
if cls._numa_node_has_cpus(numa_node):
|
|
return numa_node
|
|
# On non-CDMM Grace-Blackwell systems (e.g. GB200), each GPU's HBM
|
|
# is a separate NUMA node with no CPUs. Fall through to
|
|
# CPU-affinity-based detection to find the nearest CPU node.
|
|
logger.debug(
|
|
"NUMA node %d for GPU %d has no CPUs (non-CDMM topology), "
|
|
"falling back to CPU-affinity-based detection",
|
|
numa_node,
|
|
device_id,
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
cpu_ids = cls._get_device_cpu_affinity(handle)
|
|
if cpu_ids:
|
|
numa_node = cls._get_numa_node_for_cpu(cpu_ids[0])
|
|
if numa_node is not None:
|
|
logger.debug(
|
|
"Determined NUMA node %d for GPU %d via CPU affinity",
|
|
numa_node,
|
|
device_id,
|
|
)
|
|
return numa_node
|
|
except Exception as e:
|
|
logger.warning("Failed to get NUMA node for GPU %d: %s", device_id, e)
|
|
|
|
return None
|
|
|
|
@classmethod
|
|
def _numa_node_has_cpus(cls, node_id: int) -> bool:
|
|
"""Check whether a NUMA node has any CPUs assigned to it."""
|
|
from pathlib import Path
|
|
|
|
cpulist_file = Path(f"/sys/devices/system/node/node{node_id}/cpulist")
|
|
try:
|
|
return cpulist_file.read_text().strip() != ""
|
|
except (OSError, ValueError):
|
|
return False
|
|
|
|
@classmethod
|
|
def _get_device_cpu_affinity(cls, handle) -> list[int]:
|
|
"""Get the list of CPU IDs associated with a GPU via NVML."""
|
|
cpu_count = os.cpu_count()
|
|
if cpu_count is None:
|
|
return []
|
|
|
|
cpu_set_size = (cpu_count + 63) // 64
|
|
cpu_affinity_mask = pynvml.nvmlDeviceGetCpuAffinity(handle, cpu_set_size)
|
|
|
|
cpu_ids = []
|
|
for i, mask in enumerate(cpu_affinity_mask):
|
|
for bit in range(64):
|
|
cpu_id = i * 64 + bit
|
|
if cpu_id >= cpu_count:
|
|
break
|
|
if mask & (1 << bit):
|
|
cpu_ids.append(cpu_id)
|
|
return cpu_ids
|
|
|
|
@classmethod
|
|
def _get_numa_node_for_cpu(cls, cpu_id: int) -> int | None:
|
|
"""Determine which NUMA node a CPU belongs to."""
|
|
from pathlib import Path
|
|
|
|
node_path = Path("/sys/devices/system/node")
|
|
if not node_path.exists():
|
|
return None
|
|
|
|
for node_dir in node_path.iterdir():
|
|
if not node_dir.name.startswith("node"):
|
|
continue
|
|
try:
|
|
node_id = int(node_dir.name[4:])
|
|
cpulist_file = node_dir / "cpulist"
|
|
if cpulist_file.exists():
|
|
cpulist = cpulist_file.read_text().strip()
|
|
if cls._cpu_in_cpulist(cpu_id, cpulist):
|
|
return node_id
|
|
except (ValueError, OSError):
|
|
continue
|
|
return None
|
|
|
|
@classmethod
|
|
def _cpu_in_cpulist(cls, cpu_id: int, cpulist: str) -> bool:
|
|
"""Check if a CPU ID is in a cpulist string such as '0-3,8-11'."""
|
|
for part in cpulist.split(","):
|
|
part = part.strip()
|
|
if "-" in part:
|
|
start, end = part.split("-", 1)
|
|
if int(start) <= cpu_id <= int(end):
|
|
return True
|
|
elif part.isdigit() and int(part) == cpu_id:
|
|
return True
|
|
return False
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def get_all_device_numa_nodes(cls) -> list[int] | None:
|
|
"""Get NUMA nodes for all visible GPU devices."""
|
|
try:
|
|
numa_nodes = []
|
|
for device_id in range(cls.device_count()):
|
|
numa_node = cls.get_device_numa_node(device_id)
|
|
if numa_node is None:
|
|
logger.warning(
|
|
"Could not detect NUMA node for GPU %d, "
|
|
"disabling automatic NUMA binding",
|
|
device_id,
|
|
)
|
|
return None
|
|
numa_nodes.append(numa_node)
|
|
return numa_nodes
|
|
except Exception as e:
|
|
logger.warning("Failed to get NUMA nodes for GPUs: %s", e)
|
|
return None
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def get_all_gpu_pci_bus_ids(cls) -> dict[int, str]:
|
|
"""Query NVML for GPU index -> PCI bus ID mapping."""
|
|
out: dict[int, str] = {}
|
|
for idx in range(pynvml.nvmlDeviceGetCount()):
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
|
|
pci_info = pynvml.nvmlDeviceGetPciInfo(handle)
|
|
bus_id = pci_info.busId
|
|
if isinstance(bus_id, bytes):
|
|
bus_id = bus_id.decode("utf-8")
|
|
out[idx] = bus_id.rstrip("\x00")
|
|
if not out:
|
|
raise RuntimeError("NVML returned no GPU PCI bus ID rows")
|
|
return out
|
|
|
|
@classmethod
|
|
@with_nvml_context
|
|
def log_warnings(cls):
|
|
device_ids: int = pynvml.nvmlDeviceGetCount()
|
|
if device_ids > 1:
|
|
device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
|
|
if (
|
|
len(set(device_names)) > 1
|
|
and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"
|
|
):
|
|
logger.warning(
|
|
"Detected different devices in the system: %s. Please"
|
|
" make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
|
|
"avoid unexpected behavior.",
|
|
", ".join(device_names),
|
|
)
|
|
|
|
|
|
class NonNvmlCudaPlatform(CudaPlatformBase):
|
|
@classmethod
|
|
@cache
|
|
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
|
|
major, minor = torch.cuda.get_device_capability(device_id)
|
|
return DeviceCapability(major=major, minor=minor)
|
|
|
|
@classmethod
|
|
def get_device_name(cls, device_id: int = 0) -> str:
|
|
return torch.cuda.get_device_name(device_id)
|
|
|
|
@classmethod
|
|
def get_device_total_memory(cls, device_id: int = 0) -> int:
|
|
device_props = torch.cuda.get_device_properties(device_id)
|
|
return device_props.total_memory
|
|
|
|
@classmethod
|
|
def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
|
|
logger.exception(
|
|
"NVLink detection not possible, as context support was"
|
|
" not found. Assuming no NVLink available."
|
|
)
|
|
return False
|
|
|
|
@classmethod
|
|
def get_device_numa_node(cls, device_id: int = 0) -> int | None:
|
|
return None
|
|
|
|
@classmethod
|
|
def get_all_device_numa_nodes(cls) -> list[int] | None:
|
|
return None
|
|
|
|
|
|
# Autodetect either NVML-enabled or non-NVML platform
|
|
# based on whether NVML is available.
|
|
nvml_available = False
|
|
try:
|
|
try:
|
|
pynvml.nvmlInit()
|
|
nvml_available = True
|
|
except Exception:
|
|
# On Jetson, NVML is not supported.
|
|
nvml_available = False
|
|
finally:
|
|
if nvml_available:
|
|
pynvml.nvmlShutdown()
|
|
|
|
CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
|
|
|
|
CudaPlatform.log_warnings()
|