| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import importlib.metadata |
| 5 | import os |
| 6 | import random |
| 7 | import sys |
| 8 | import threading |
| 9 | from collections.abc import Callable, Collection |
| 10 | from typing import TYPE_CHECKING, Any, TypeVar |
| 11 | |
| 12 | import numpy as np |
| 13 | import numpy.typing as npt |
| 14 | import torch |
| 15 | from packaging import version |
| 16 | from packaging.version import Version |
| 17 | from torch.library import Library, infer_schema |
| 18 | |
| 19 | import vllm.envs as envs |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.utils.platform_utils import is_pin_memory_available |
| 22 | |
| 23 | if TYPE_CHECKING: |
| 24 | from vllm.config import ModelConfig |
| 25 | from vllm.sequence import IntermediateTensors |
| 26 | else: |
| 27 | ModelConfig = object |
| 28 | IntermediateTensors = object |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | STR_DTYPE_TO_TORCH_DTYPE = { |
| 34 | "float32": torch.float32, |
| 35 | "half": torch.half, |
| 36 | "float16": torch.float16, |
| 37 | "bfloat16": torch.bfloat16, |
| 38 | "float": torch.float, |
| 39 | "fp8": torch.uint8, |
| 40 | "fp8_e4m3": torch.uint8, |
| 41 | "fp8_e5m2": torch.uint8, |
| 42 | "int8": torch.int8, |
| 43 | "int4_per_token_head": torch.uint8, |
| 44 | "int8_per_token_head": torch.int8, |
| 45 | "fp8_per_token_head": torch.uint8, |
| 46 | "fp8_inc": torch.float8_e4m3fn, |
| 47 | "fp8_ds_mla": torch.uint8, |
| 48 | "nvfp4_ds_mla": torch.uint8, |
| 49 | "turboquant_k8v4": torch.uint8, |
| 50 | "turboquant_4bit_nc": torch.uint8, |
| 51 | "turboquant_k3v4_nc": torch.uint8, |
| 52 | "turboquant_3bit_nc": torch.uint8, |
| 53 | "nvfp4": torch.uint8, |
| 54 | "nvfp4_4over6": torch.uint8, |
| 55 | } |
| 56 | |
| 57 | TORCH_DTYPE_TO_NUMPY_DTYPE = { |
| 58 | torch.float16: np.float16, |
| 59 | torch.float32: np.float32, |
| 60 | torch.float64: np.float64, |
| 61 | torch.uint8: np.uint8, |
| 62 | torch.int32: np.int32, |
| 63 | torch.int64: np.int64, |
| 64 | } |
| 65 | |
| 66 | |
| 67 | MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = { |
| 68 | "fp8": "fp8_e4m3", |
| 69 | "nvfp4": "nvfp4", |
| 70 | } |
| 71 | |
| 72 | T = TypeVar("T") |
| 73 | |
| 74 | |
| 75 | PIN_MEMORY = is_pin_memory_available() |
| 76 | |
| 77 | |
| 78 | def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: |
| 79 | return ( |
| 80 | kv_cache_dtype.startswith("fp8") |
| 81 | or kv_cache_dtype.endswith("per_token_head") |
| 82 | or kv_cache_dtype.startswith("nvfp4") |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: |
| 87 | """Return True if *kv_cache_dtype* needs per-token-head scales.""" |
| 88 | return kv_cache_dtype.endswith("per_token_head") |
| 89 | |
| 90 | |
| 91 | def is_meta_module(module: torch.nn.Module) -> bool: |
| 92 | """Return True if module contains any meta parameters or buffers.""" |
| 93 | return any(p.is_meta for p in module.parameters()) or any( |
| 94 | b.is_meta for b in module.buffers() |
| 95 | ) |
| 96 | |
| 97 | |
| 98 | def is_strictly_contiguous(t: torch.Tensor) -> bool: |
| 99 | """Check if tensor is contiguous AND has no degenerate strides. |
| 100 | |
| 101 | A degenerate stride occurs when a dimension has size 1 but the stride |
| 102 | doesn't match the canonical contiguous layout. This can cause issues |
| 103 | in some CUDA kernels that rely on stride values for memory access. |
| 104 | |
| 105 | For a C-contiguous tensor of shape (d0, d1, ..., dn), the expected |
| 106 | strides are: stride[i] = product(shape[i+1:]) for all i, with stride[-1]=1. |
| 107 | |
| 108 | Example with torch.Size([16, 1, 8, 32]): |
| 109 | - Canonical strides: (256, 256, 32, 1) |
| 110 | - Degenerate strides: (256, 1, 32, 1) # dim=1 has size=1, allowing |
| 111 | # non-canonical stride in dim=0 |
| 112 | """ |
| 113 | if not t.is_contiguous(): |
| 114 | return False |
| 115 | |
| 116 | # Check that strides match canonical contiguous layout |
| 117 | shape = t.shape |
| 118 | strides = t.stride() |
| 119 | expected_stride = 1 |
| 120 | for i in range(len(shape) - 1, -1, -1): |
| 121 | if strides[i] != expected_stride: |
| 122 | return False |
| 123 | expected_stride *= shape[i] |
| 124 | return True |
| 125 | |
| 126 | |
| 127 | def is_non_overlapping_and_dense(t: torch.Tensor) -> bool: |
| 128 | """Check if the tensor's elements cover one gapless, non-overlapping byte |
| 129 | range, in any dimension order (i.e. a permuted view of a contiguous |
| 130 | tensor); ``is_contiguous()`` additionally requires row-major order. |
| 131 | """ |
| 132 | expected_stride = 1 |
| 133 | for size, stride in sorted(zip(t.shape, t.stride()), key=lambda p: p[1]): |
| 134 | if size == 1: |
| 135 | continue |
| 136 | if stride != expected_stride: |
| 137 | return False |
| 138 | expected_stride *= size |
| 139 | return True |
| 140 | |
| 141 | |
| 142 | def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor: |
| 143 | """Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility. |
| 144 | |
| 145 | PyTorch allows any stride on a size=1 dim (is_contiguous() is always True |
| 146 | there), so a size=1 dim may have stride=1 (2 bytes for bf16) instead of |
| 147 | the canonical product(shape[i+1:]). CUDA TMA on H100+ requires all |
| 148 | non-outermost strides to be ≥16-byte aligned; stride=1 triggers |
| 149 | cudaErrorIllegalInstruction. Zero-copy: patches stride metadata only via |
| 150 | as_strided; returns t unchanged if all size=1 strides are already canonical. |
| 151 | """ |
| 152 | if 1 not in t.shape: |
| 153 | return t |
| 154 | strides = list(t.stride()) |
| 155 | shape = t.shape |
| 156 | prev_stride = 1 |
| 157 | changed = False |
| 158 | for i in range(len(shape) - 1, -1, -1): |
| 159 | if shape[i] == 1 and strides[i] != prev_stride: |
| 160 | strides[i] = prev_stride |
| 161 | changed = True |
| 162 | prev_stride = strides[i] * shape[i] |
| 163 | if not changed: |
| 164 | return t |
| 165 | return t.as_strided(t.shape, strides) |
| 166 | |
| 167 | |
| 168 | @contextlib.contextmanager |
| 169 | def set_default_torch_dtype(dtype: torch.dtype): |
| 170 | """Sets the default torch dtype to the given dtype.""" |
| 171 | old_dtype = torch.get_default_dtype() |
| 172 | torch.set_default_dtype(dtype) |
| 173 | try: |
| 174 | yield |
| 175 | finally: |
| 176 | torch.set_default_dtype(old_dtype) |
| 177 | |
| 178 | |
| 179 | def _cgroup_cpu_limit() -> float | None: |
| 180 | """Effective CPU quota of this process's cgroup, None if unlimited. |
| 181 | |
| 182 | Resolves the process's own cgroup from /proc/self/cgroup and takes the |
| 183 | tightest cpu.max (v2) or cfs quota (v1) along the hierarchy. |
| 184 | """ |
| 185 | limit: float | None = None |
| 186 | try: |
| 187 | with open("/proc/self/cgroup") as f: |
| 188 | entries = [line.strip().split(":", 2) for line in f] |
| 189 | |
| 190 | def visit(base: str, rel_path: str, read_quota) -> None: |
| 191 | nonlocal limit |
| 192 | path = rel_path |
| 193 | while path: |
| 194 | quota = read_quota(os.path.join(base, path.lstrip("/"))) |
| 195 | if quota is not None: |
| 196 | limit = quota if limit is None else min(limit, quota) |
| 197 | path = path.rsplit("/", 1)[0] |
| 198 | |
| 199 | def read_v2(cg_dir: str) -> float | None: |
| 200 | try: |
| 201 | with open(os.path.join(cg_dir, "cpu.max")) as f: |
| 202 | quota, period = f.read().split() |
| 203 | return None if quota == "max" else float(quota) / float(period) |
| 204 | except (OSError, ValueError): |
| 205 | return None |
| 206 | |
| 207 | def read_v1(cg_dir: str) -> float | None: |
| 208 | try: |
| 209 | with open(os.path.join(cg_dir, "cpu.cfs_quota_us")) as f: |
| 210 | quota = int(f.read()) |
| 211 | if quota <= 0: |
| 212 | return None |
| 213 | with open(os.path.join(cg_dir, "cpu.cfs_period_us")) as f: |
| 214 | return quota / int(f.read()) |
| 215 | except (OSError, ValueError): |
| 216 | return None |
| 217 | |
| 218 | for entry in entries: |
| 219 | if len(entry) != 3: |
| 220 | continue |
| 221 | _, controllers, rel_path = entry |
| 222 | if controllers == "": # cgroup v2 |
| 223 | visit("/sys/fs/cgroup", rel_path, read_v2) |
| 224 | elif "cpu" in controllers.split(","): # cgroup v1 |
| 225 | visit("/sys/fs/cgroup/cpu", rel_path, read_v1) |
| 226 | except OSError: |
| 227 | pass |
| 228 | return limit |
| 229 | |
| 230 | |
| 231 | def available_cpu_count() -> int: |
| 232 | """CPUs actually usable by this process: scheduling affinity capped by |
| 233 | the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind).""" |
| 234 | if sys.platform != "linux": |
| 235 | return os.cpu_count() or 1 |
| 236 | count = len(os.sched_getaffinity(0)) |
| 237 | limit = _cgroup_cpu_limit() |
| 238 | if limit is not None: |
| 239 | count = min(count, int(limit)) |
| 240 | return max(1, count) |
| 241 | |
| 242 | |
| 243 | # Marks OMP_NUM_THREADS as chosen by vLLM for its worker processes rather than |
| 244 | # set by the user, so a worker knows it may drop the value once startup is done. |
| 245 | OMP_NUM_THREADS_SET_BY_VLLM = "VLLM_OMP_NUM_THREADS_SET_BY_VLLM" |
| 246 | |
| 247 | |
| 248 | def startup_omp_num_threads(num_local_procs: int) -> int: |
| 249 | """Thread count for a worker process's startup work (weight loading). |
| 250 | |
| 251 | Weight loading does CPU-parallel work, so workers benefit from more than |
| 252 | one thread, but only a bounded share of the CPUs this node's workers may |
| 253 | actually use: torch's default is the host core count, which ignores both |
| 254 | scheduling affinity and any cgroup CPU quota, and doesn't account for the |
| 255 | other workers sharing the node. |
| 256 | """ |
| 257 | return max(1, available_cpu_count() // max(1, num_local_procs)) |
| 258 | |
| 259 | |
| 260 | def set_torch_threads_for_runtime() -> None: |
| 261 | """Set torch intra-op threads to 1 for steady-state serving. |
| 262 | |
| 263 | Any multi-threaded torch CPU op in the engine hot loop leaves the OMP |
| 264 | workers spin-waiting after each parallel region, stealing cycles from the |
| 265 | step's serial code (and burning cgroup CPU quota in containers). No |
| 266 | steady-state CPU op benefits from intra-op parallelism. |
| 267 | Respects an externally-set OMP_NUM_THREADS. |
| 268 | """ |
| 269 | if ( |
| 270 | omp_num_threads := os.environ.get("OMP_NUM_THREADS") |
| 271 | ) is not None and os.environ.get(OMP_NUM_THREADS_SET_BY_VLLM) != "1": |
| 272 | try: |
| 273 | if int(omp_num_threads) > 1: |
| 274 | logger.warning_once( |
| 275 | "OMP_NUM_THREADS=%s is set; leaving Torch threads at %d " |
| 276 | "for serving. Multi-threaded torch CPU ops during serving " |
| 277 | "can degrade performance through spin-wait contention and " |
| 278 | "cgroup CPU-quota throttling.", |
| 279 | omp_num_threads, |
| 280 | torch.get_num_threads(), |
| 281 | ) |
| 282 | except ValueError: |
| 283 | pass |
| 284 | return |
| 285 | if torch.get_num_threads() != 1: |
| 286 | logger.info_once( |
| 287 | "Reducing Torch threads from %d to 1 for serving. Set " |
| 288 | "OMP_NUM_THREADS in the external environment to override.", |
| 289 | torch.get_num_threads(), |
| 290 | ) |
| 291 | torch.set_num_threads(1) |
| 292 | |
| 293 | |
| 294 | @contextlib.contextmanager |
| 295 | def set_default_torch_num_threads(num_threads: int | None = None): |
| 296 | """Sets the default number of threads for PyTorch to the given value. |
| 297 | |
| 298 | `None` means using the value of the environment variable `OMP_NUM_THREADS` |
| 299 | (or `1` if that is not available). |
| 300 | """ |
| 301 | if num_threads is None: |
| 302 | num_threads = 1 |
| 303 | |
| 304 | try: |
| 305 | num_threads = int(os.environ["OMP_NUM_THREADS"]) |
| 306 | except KeyError: |
| 307 | logger.debug_once( |
| 308 | "OMP_NUM_THREADS is not set; defaulting Torch threads to %d.", |
| 309 | num_threads, |
| 310 | ) |
| 311 | except ValueError: |
| 312 | logger.warning_once( |
| 313 | "OMP_NUM_THREADS is invalid; defaulting Torch threads to %d.", |
| 314 | num_threads, |
| 315 | ) |
| 316 | |
| 317 | old_num_threads = torch.get_num_threads() |
| 318 | torch.set_num_threads(num_threads) |
| 319 | |
| 320 | try: |
| 321 | yield |
| 322 | finally: |
| 323 | torch.set_num_threads(old_num_threads) |
| 324 | |
| 325 | |
| 326 | @contextlib.contextmanager |
| 327 | def guard_cuda_initialization(): |
| 328 | """Avoid unexpected CUDA initialization.""" |
| 329 | from vllm.platforms import current_platform |
| 330 | |
| 331 | if not current_platform.is_cuda(): |
| 332 | yield |
| 333 | return |
| 334 | |
| 335 | old_value = os.environ.get("CUDA_VISIBLE_DEVICES") |
| 336 | os.environ["CUDA_VISIBLE_DEVICES"] = "" |
| 337 | try: |
| 338 | yield |
| 339 | except Exception as e: |
| 340 | if "No CUDA GPUs are available" in str(e): |
| 341 | err_msg = "CUDA initialization is blocked." |
| 342 | else: |
| 343 | err_msg = str(e) |
| 344 | raise RuntimeError(err_msg) from e |
| 345 | finally: |
| 346 | if old_value is None: |
| 347 | del os.environ["CUDA_VISIBLE_DEVICES"] |
| 348 | else: |
| 349 | os.environ["CUDA_VISIBLE_DEVICES"] = old_value |
| 350 | |
| 351 | |
| 352 | def get_dtype_size(dtype: torch.dtype) -> int: |
| 353 | """Get the size of the data type in bytes.""" |
| 354 | return torch.tensor([], dtype=dtype).element_size() |
| 355 | |
| 356 | |
| 357 | # bool = 0, int = 1, float = 2, complex = 3 |
| 358 | def _get_precision_level(dtype: torch.dtype) -> int: |
| 359 | # NOTE: Complex dtypes return `is_floating_point=False` |
| 360 | return (dtype != torch.bool) + dtype.is_floating_point + dtype.is_complex * 2 |
| 361 | |
| 362 | |
| 363 | def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype): |
| 364 | """Test whether it is lossless to cast a tensor from |
| 365 | `src_dtype` to `tgt_dtype`. |
| 366 | """ |
| 367 | if src_dtype == tgt_dtype: |
| 368 | return True |
| 369 | |
| 370 | src_level = _get_precision_level(src_dtype) |
| 371 | tgt_level = _get_precision_level(tgt_dtype) |
| 372 | |
| 373 | if src_level < tgt_level: |
| 374 | return True |
| 375 | if src_level > tgt_level: |
| 376 | return False |
| 377 | |
| 378 | # Compare integral types |
| 379 | if not src_dtype.is_floating_point and not src_dtype.is_complex: |
| 380 | src_info = torch.iinfo(src_dtype) |
| 381 | tgt_info = torch.iinfo(tgt_dtype) |
| 382 | return src_info.min >= tgt_info.min and src_info.max <= tgt_info.max |
| 383 | |
| 384 | # Compare floating-point types |
| 385 | src_info = torch.finfo(src_dtype) |
| 386 | tgt_info = torch.finfo(tgt_dtype) |
| 387 | return ( |
| 388 | src_info.min >= tgt_info.min |
| 389 | and src_info.max <= tgt_info.max |
| 390 | and src_info.resolution >= tgt_info.resolution |
| 391 | ) |
| 392 | |
| 393 | |
| 394 | def common_broadcastable_dtype(dtypes: Collection[torch.dtype]): |
| 395 | """Get the common `dtype` where all of the other `dtypes` can be |
| 396 | cast to it without losing any information. |
| 397 | """ |
| 398 | return max( |
| 399 | dtypes, |
| 400 | key=lambda dtype: sum(is_lossless_cast(dt, dtype) for dt in dtypes), |
| 401 | ) |
| 402 | |
| 403 | |
| 404 | def _generate_random_fp8( |
| 405 | tensor: torch.Tensor, |
| 406 | low: float, |
| 407 | high: float, |
| 408 | ) -> None: |
| 409 | # NOTE(zhaoyang): Due to NaN and Inf representation for fp8 data type, |
| 410 | # it may occur Inf or NaN if we directly use torch.randint |
| 411 | # to generate random data for fp8 data. |
| 412 | # For example, s.11111.00 in fp8e5m2 format represents Inf. |
| 413 | # | E4M3 | E5M2 |
| 414 | # -----|-------------|------------------- |
| 415 | # Inf | N/A | s.11111.00 |
| 416 | # NaN | s.1111.111 | s.11111.{01,10,11} |
| 417 | from vllm import _custom_ops as ops |
| 418 | |
| 419 | tensor_tmp = torch.empty_like(tensor, dtype=torch.float16) |
| 420 | tensor_tmp.uniform_(low, high) |
| 421 | ops.convert_fp8(tensor, tensor_tmp) |
| 422 | del tensor_tmp |
| 423 | |
| 424 | |
| 425 | def get_kv_cache_torch_dtype( |
| 426 | cache_dtype: str | torch.dtype | None, |
| 427 | model_dtype: str | torch.dtype | None = None, |
| 428 | ) -> torch.dtype: |
| 429 | if isinstance(cache_dtype, str): |
| 430 | if cache_dtype == "auto": |
| 431 | if isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE: |
| 432 | torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype] |
| 433 | elif isinstance(model_dtype, torch.dtype): |
| 434 | torch_dtype = model_dtype |
| 435 | else: |
| 436 | raise ValueError(f"Invalid model dtype: {model_dtype}") |
| 437 | elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE: |
| 438 | torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype] |
| 439 | else: |
| 440 | raise ValueError(f"Invalid kv cache dtype: {cache_dtype}") |
| 441 | elif isinstance(cache_dtype, torch.dtype): |
| 442 | torch_dtype = cache_dtype |
| 443 | else: |
| 444 | raise ValueError(f"Invalid kv cache dtype: {cache_dtype}") |
| 445 | return torch_dtype |
| 446 | |
| 447 | |
| 448 | def get_kv_cache_quant_algo_string(quant_cfg: dict[str, Any]) -> str | None: |
| 449 | """Get the KV cache quantization algorithm string from the quantization config. |
| 450 | |
| 451 | Maps various FP8 format names to vLLM's standard cache dtype strings. |
| 452 | Returns None if no kv_cache_quant_algo is specified. |
| 453 | Returns "auto" if the value is not recognized/supported. |
| 454 | """ |
| 455 | # Mapping from model config values to vLLM cache_dtype strings |
| 456 | |
| 457 | quant_method = quant_cfg.get("quant_method", "") |
| 458 | if quant_method.startswith("modelopt"): |
| 459 | quantization_inner = quant_cfg.get("quantization", quant_cfg) |
| 460 | # Check if quant config is specified and use kv cache quant algo |
| 461 | kv_algo = ( |
| 462 | quantization_inner.get("kv_cache_scheme") |
| 463 | or quant_cfg.get("kv_cache_scheme") |
| 464 | or quantization_inner.get("kv_cache_quant_algo") |
| 465 | or quant_cfg.get("kv_cache_quant_algo") |
| 466 | ) |
| 467 | if isinstance(kv_algo, dict): |
| 468 | if ( |
| 469 | kv_algo.get("dynamic") is False |
| 470 | and kv_algo.get("num_bits") == 8 |
| 471 | and kv_algo.get("type") == "float" |
| 472 | ): |
| 473 | kv_algo = "fp8" |
| 474 | elif kv_algo.get("num_bits") == 4 and kv_algo.get("type") == "float": |
| 475 | kv_algo = "nvfp4" |
| 476 | else: |
| 477 | # Unknown/unsupported format - return "auto" as safe fallback |
| 478 | logger.warning( |
| 479 | "WARNING: Unknown kv_cache_quant_algo '%s' in model " |
| 480 | "config. Supported values: %s. Falling back to 'auto'.", |
| 481 | f"{kv_algo}", |
| 482 | list(MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP.keys()), |
| 483 | ) |
| 484 | return "auto" |
| 485 | if isinstance(kv_algo, str): |
| 486 | kv_algo_lower = kv_algo.lower() |
| 487 | |
| 488 | # Try to map to vLLM's standard format |
| 489 | if kv_algo_lower in MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP: |
| 490 | return MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP[kv_algo_lower] |
| 491 | else: |
| 492 | # Unknown/unsupported format - return "auto" as safe fallback |
| 493 | logger.warning( |
| 494 | "WARNING: Unknown kv_cache_quant_algo '%s' in model " |
| 495 | "config. Supported values: %s. Falling back to 'auto'.", |
| 496 | kv_algo, |
| 497 | list(MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP.keys()), |
| 498 | ) |
| 499 | return "auto" |
| 500 | return None |
| 501 | |
| 502 | |
| 503 | def resolve_kv_cache_dtype_string( |
| 504 | kv_cache_dtype: str, model_config: ModelConfig |
| 505 | ) -> str: |
| 506 | """Resolve 'auto' kv_cache_dtype to the actual string value from model config. |
| 507 | Returns the resolved cache_dtype string. |
| 508 | """ |
| 509 | if kv_cache_dtype != "auto": |
| 510 | return kv_cache_dtype |
| 511 | |
| 512 | hf_cfg = getattr(model_config, "hf_config", None) |
| 513 | if hf_cfg is not None: |
| 514 | quant_cfg = getattr(hf_cfg, "quantization_config", None) |
| 515 | if quant_cfg is not None: |
| 516 | kv_algo_str = get_kv_cache_quant_algo_string(quant_cfg) |
| 517 | if kv_algo_str is not None: |
| 518 | return kv_algo_str |
| 519 | |
| 520 | # Default to auto (will be handled by downstream code) |
| 521 | return "auto" |
| 522 | |
| 523 | |
| 524 | def kv_cache_dtype_str_to_dtype( |
| 525 | kv_cache_dtype: str, model_config: ModelConfig |
| 526 | ) -> torch.dtype: |
| 527 | if kv_cache_dtype == "auto": |
| 528 | # Model config may not be specified for unit tests, default to float16 |
| 529 | return model_config.dtype if model_config else torch.half |
| 530 | return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype] |
| 531 | |
| 532 | |
| 533 | def set_random_seed(seed: int | None) -> None: |
| 534 | if seed is not None: |
| 535 | random.seed(seed) |
| 536 | np.random.seed(seed) |
| 537 | torch.manual_seed(seed) |
| 538 | from vllm.platforms import current_platform |
| 539 | |
| 540 | current_platform.manual_seed_all(seed) |
| 541 | |
| 542 | |
| 543 | def nvfp4_kv_cache_full_dim(head_size: int) -> int: |
| 544 | """Packed last dim for NVFP4 KV cache: fp4 data + fp8 block scales.""" |
| 545 | return head_size // 2 + head_size // 16 |
| 546 | |
| 547 | |
| 548 | def nvfp4_split_data_scale( |
| 549 | kv_side: torch.Tensor, |
| 550 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 551 | """Split one side (K or V) of an NVFP4 KV cache into data and scale. |
| 552 | |
| 553 | The input is a 4D uint8 tensor whose last dimension is |
| 554 | ``full_dim = data_dim + scale_dim``. The physical layout within each |
| 555 | side is ``[data | scale]``, both packed contiguously. |
| 556 | |
| 557 | The caller is responsible for slicing K and V from the combined cache |
| 558 | first (e.g. ``kv_cache.split(num_kv_heads, dim=1)``). |
| 559 | |
| 560 | Args: |
| 561 | kv_side: 4D uint8 tensor ``(B, H, N, full_dim)``. |
| 562 | |
| 563 | Returns: |
| 564 | ``(data, scale)`` where *data* is uint8 and *scale* is |
| 565 | float8_e4m3fn, both views of the same storage. |
| 566 | |
| 567 | """ |
| 568 | num_pages = kv_side.shape[0] |
| 569 | dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2] |
| 570 | full_dim = kv_side.shape[3] |
| 571 | data_dim = full_dim * 8 // 9 |
| 572 | scale_dim = full_dim - data_dim |
| 573 | |
| 574 | data_per_kv = dim_1 * dim_2 * data_dim |
| 575 | page_bytes = kv_side.stride(0) |
| 576 | |
| 577 | # Derive inner strides from the kv_side strides, scaling by the |
| 578 | # ratio of the target dim to full_dim. This preserves the physical |
| 579 | # layout (NHD vs HND) encoded in the input tensor's strides. |
| 580 | s1 = kv_side.stride(1) * data_dim // full_dim |
| 581 | s2 = kv_side.stride(2) * data_dim // full_dim |
| 582 | data_shape = (num_pages, dim_1, dim_2, data_dim) |
| 583 | data_strides = (page_bytes, s1, s2, 1) |
| 584 | |
| 585 | s1_s = kv_side.stride(1) * scale_dim // full_dim |
| 586 | s2_s = kv_side.stride(2) * scale_dim // full_dim |
| 587 | scale_shape = (num_pages, dim_1, dim_2, scale_dim) |
| 588 | scale_strides = (page_bytes, s1_s, s2_s, 1) |
| 589 | |
| 590 | base = kv_side.storage_offset() |
| 591 | data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base) |
| 592 | scale = torch.as_strided( |
| 593 | kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv |
| 594 | ).view(torch.float8_e4m3fn) |
| 595 | |
| 596 | return data, scale |
| 597 | |
| 598 | |
| 599 | def create_kv_caches_with_random_flash( |
| 600 | num_blocks: int, |
| 601 | block_size: int, |
| 602 | num_layers: int, |
| 603 | num_heads: int, |
| 604 | head_size: int, |
| 605 | cache_dtype: str | torch.dtype | None, |
| 606 | model_dtype: str | torch.dtype | None = None, |
| 607 | seed: int | None = None, |
| 608 | device: str | None = "cuda", |
| 609 | cache_layout: str | None = "NHD", |
| 610 | ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: |
| 611 | set_random_seed(seed) |
| 612 | |
| 613 | dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype) |
| 614 | generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size) |
| 615 | assert cache_layout in ("NHD", "HND") |
| 616 | stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4) |
| 617 | |
| 618 | kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order) |
| 619 | scale = head_size**-0.5 |
| 620 | |
| 621 | key_caches: list[torch.Tensor] = [] |
| 622 | value_caches: list[torch.Tensor] = [] |
| 623 | |
| 624 | for _ in range(num_layers): |
| 625 | if isinstance(cache_dtype, str) and cache_dtype.startswith("nvfp4"): |
| 626 | # Full page dim: fp4 data + fp8 block scales per head. |
| 627 | # Per page layout: [K_data | K_scale | V_data | V_scale] |
| 628 | # Returns [:, 0] and [:, 1] like all other dtypes. |
| 629 | full_dim = nvfp4_kv_cache_full_dim(head_size) |
| 630 | nvfp4_shape = (num_blocks, 2, block_size, num_heads, full_dim) |
| 631 | nvfp4_phys = tuple(nvfp4_shape[i] for i in stride_order) |
| 632 | inv = [stride_order.index(i) for i in range(len(stride_order))] |
| 633 | key_value_cache = torch.randint( |
| 634 | 0, |
| 635 | 256, |
| 636 | nvfp4_phys, |
| 637 | dtype=dtype, |
| 638 | device=device, |
| 639 | ).permute(*inv) |
| 640 | else: |
| 641 | key_value_cache = torch.empty( |
| 642 | size=kv_cache_allocation_shape, dtype=dtype, device=device |
| 643 | ).permute(*stride_order) |
| 644 | if cache_dtype in ["auto", "half", "bfloat16", "float"]: |
| 645 | key_value_cache.uniform_(-scale, scale) |
| 646 | elif cache_dtype == "fp8": |
| 647 | _generate_random_fp8(key_value_cache, -scale, scale) |
| 648 | else: |
| 649 | raise ValueError(f"Does not support key cache of type {cache_dtype}") |
| 650 | key_caches.append(key_value_cache[:, 0]) |
| 651 | value_caches.append(key_value_cache[:, 1]) |
| 652 | return key_caches, value_caches |
| 653 | |
| 654 | |
| 655 | def create_kv_caches_with_random( |
| 656 | num_blocks: int, |
| 657 | block_size: int, |
| 658 | num_layers: int, |
| 659 | num_heads: int, |
| 660 | head_size: int, |
| 661 | cache_dtype: str | torch.dtype | None, |
| 662 | model_dtype: str | torch.dtype | None = None, |
| 663 | seed: int | None = None, |
| 664 | device: str | None = "cuda", |
| 665 | ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: |
| 666 | if cache_dtype == "fp8" and head_size % 16: |
| 667 | raise ValueError( |
| 668 | f"Does not support key cache of type fp8 with head_size {head_size}" |
| 669 | ) |
| 670 | |
| 671 | set_random_seed(seed) |
| 672 | |
| 673 | dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype) |
| 674 | |
| 675 | scale = head_size**-0.5 |
| 676 | x = 16 // torch.tensor([], dtype=dtype).element_size() |
| 677 | key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x) |
| 678 | key_caches: list[torch.Tensor] = [] |
| 679 | for _ in range(num_layers): |
| 680 | key_cache = torch.empty(size=key_cache_shape, dtype=dtype, device=device) |
| 681 | if cache_dtype in ["auto", "half", "bfloat16", "float"]: |
| 682 | key_cache.uniform_(-scale, scale) |
| 683 | elif cache_dtype == "fp8": |
| 684 | _generate_random_fp8(key_cache, -scale, scale) |
| 685 | else: |
| 686 | raise ValueError(f"Does not support key cache of type {cache_dtype}") |
| 687 | key_caches.append(key_cache) |
| 688 | |
| 689 | value_cache_shape = (num_blocks, num_heads, head_size, block_size) |
| 690 | value_caches: list[torch.Tensor] = [] |
| 691 | for _ in range(num_layers): |
| 692 | value_cache = torch.empty(size=value_cache_shape, dtype=dtype, device=device) |
| 693 | if cache_dtype in ["auto", "half", "bfloat16", "float"]: |
| 694 | value_cache.uniform_(-scale, scale) |
| 695 | elif cache_dtype == "fp8": |
| 696 | _generate_random_fp8(value_cache, -scale, scale) |
| 697 | else: |
| 698 | raise ValueError(f"Does not support value cache of type {cache_dtype}") |
| 699 | value_caches.append(value_cache) |
| 700 | return key_caches, value_caches |
| 701 | |
| 702 | |
| 703 | def async_tensor_h2d( |
| 704 | data: list | np.ndarray | torch.Tensor, |
| 705 | device: str | torch.device | None = None, |
| 706 | dtype: torch.dtype | None = None, |
| 707 | out: torch.Tensor | None = None, |
| 708 | ) -> torch.Tensor: |
| 709 | """Copy list/numpy array/tensor async from host to device.""" |
| 710 | if dtype is None and out is not None: |
| 711 | dtype = out.dtype |
| 712 | if isinstance(data, np.ndarray): |
| 713 | data = torch.from_numpy(data) |
| 714 | if isinstance(data, torch.Tensor): |
| 715 | t = data |
| 716 | if PIN_MEMORY and not t.is_pinned(): |
| 717 | # Stage in pinned, contiguous buffer to ensure fully async copy. |
| 718 | t = torch.empty( |
| 719 | t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True |
| 720 | ).copy_(t) |
| 721 | else: |
| 722 | t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") |
| 723 | assert t.is_cpu |
| 724 | |
| 725 | if out is not None: |
| 726 | assert out.dtype == dtype |
| 727 | return out.copy_(t, non_blocking=True) |
| 728 | assert device is not None, "must provide destination tensor or device" |
| 729 | return t.to(device=device, dtype=dtype, non_blocking=True) |
| 730 | |
| 731 | |
| 732 | def np_to_pinned_tensor(array: np.ndarray) -> torch.Tensor: |
| 733 | t = torch.from_numpy(array) |
| 734 | return t.pin_memory() if PIN_MEMORY else t |
| 735 | |
| 736 | |
| 737 | def make_ndarray_with_pad( |
| 738 | x: list[list[T]], |
| 739 | pad: T, |
| 740 | dtype: npt.DTypeLike, |
| 741 | *, |
| 742 | max_len: int | None = None, |
| 743 | ) -> npt.NDArray: |
| 744 | """Make a padded array from 2D inputs. |
| 745 | |
| 746 | The padding is applied to the end of each inner list until it reaches |
| 747 | `max_len`. |
| 748 | """ |
| 749 | if max_len is None: |
| 750 | # Unlike for most functions, map is faster than a genexpr over `len` |
| 751 | max_len = max(map(len, x), default=0) |
| 752 | |
| 753 | padded_x = np.full((len(x), max_len), pad, dtype=dtype) |
| 754 | for ind, blocktb in enumerate(x): |
| 755 | assert len(blocktb) <= max_len |
| 756 | padded_x[ind, : len(blocktb)] = blocktb |
| 757 | |
| 758 | return padded_x |
| 759 | |
| 760 | |
| 761 | def make_tensor_with_pad( |
| 762 | x: list[list[T]], |
| 763 | pad: T, |
| 764 | dtype: torch.dtype, |
| 765 | *, |
| 766 | max_len: int | None = None, |
| 767 | device: str | torch.device | None = None, |
| 768 | pin_memory: bool = False, |
| 769 | ) -> torch.Tensor: |
| 770 | """Make a padded tensor from 2D inputs. |
| 771 | |
| 772 | The padding is applied to the end of each inner list until it reaches |
| 773 | `max_len`. |
| 774 | """ |
| 775 | np_dtype = TORCH_DTYPE_TO_NUMPY_DTYPE[dtype] |
| 776 | padded_x = make_ndarray_with_pad(x, pad, np_dtype, max_len=max_len) |
| 777 | |
| 778 | tensor = torch.from_numpy(padded_x).to(device) |
| 779 | if pin_memory: |
| 780 | tensor = tensor.pin_memory() |
| 781 | |
| 782 | return tensor |
| 783 | |
| 784 | |
| 785 | prev_set_stream = torch.cuda.set_stream |
| 786 | |
| 787 | _current_stream_tls = threading.local() |
| 788 | |
| 789 | |
| 790 | def _patched_set_stream(stream: torch.cuda.Stream) -> None: |
| 791 | _current_stream_tls.value = stream |
| 792 | prev_set_stream(stream) |
| 793 | |
| 794 | |
| 795 | torch.cuda.set_stream = _patched_set_stream |
| 796 | |
| 797 | |
| 798 | class _StreamPlaceholder: |
| 799 | def __init__(self): |
| 800 | self.synchronize = lambda: None |
| 801 | |
| 802 | |
| 803 | def current_stream() -> torch.cuda.Stream: |
| 804 | """Replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`. |
| 805 | it turns out that `torch.cuda.current_stream()` is quite expensive, |
| 806 | as it will construct a new stream object at each call. |
| 807 | here we patch `torch.cuda.set_stream` to keep track of the current stream |
| 808 | directly, so that we can avoid calling `torch.cuda.current_stream()`. |
| 809 | |
| 810 | the underlying hypothesis is that we do not call `torch._C._cuda_setStream` |
| 811 | from C/C++ code. |
| 812 | """ |
| 813 | from vllm.platforms import current_platform |
| 814 | |
| 815 | if not hasattr(_current_stream_tls, "value") or _current_stream_tls.value is None: |
| 816 | # when this function is called before any stream is set, |
| 817 | # we return the default stream. |
| 818 | # On ROCm using the default 0 stream in combination with RCCL |
| 819 | # is hurting performance. |
| 820 | # On CUDA, we capture and replay cudagraph on the same stream, |
| 821 | # so we need to avoid using the default stream as well. The default |
| 822 | # stream cannot be used for cudagraph capture, see |
| 823 | # https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77 |
| 824 | # for more details. Therefore, we create a dedicated stream per process. |
| 825 | if current_platform.is_rocm() or current_platform.is_cuda(): |
| 826 | # Ensure new stream is ordered w.r.t. replaced stream's work. |
| 827 | new_stream = torch.cuda.Stream() |
| 828 | new_stream.wait_stream(torch.cuda.current_stream()) |
| 829 | # torch.cuda.set_stream here is the alias of _pathed_set_stream |
| 830 | torch.cuda.set_stream(new_stream) |
| 831 | elif current_platform.is_cpu(): |
| 832 | _current_stream_tls.value = _StreamPlaceholder() |
| 833 | else: |
| 834 | current_stream = current_platform.current_stream |
| 835 | if current_stream is not None: |
| 836 | _current_stream_tls.value = current_stream() |
| 837 | else: |
| 838 | raise ValueError( |
| 839 | "Fail to set current stream, current platform " |
| 840 | "may not support current_stream with torch API" |
| 841 | ) |
| 842 | return _current_stream_tls.value |
| 843 | |
| 844 | |
| 845 | # Global auxiliary stream for running operations in background streams. |
| 846 | # We have single global auxiliary stream to avoid an explosion of streams |
| 847 | # for every layer (and make profiling look sane). |
| 848 | # |
| 849 | # aux_stream() is currently used for: |
| 850 | # - MoE shared_expert overlap with router |
| 851 | _aux_stream: torch.cuda.Stream | None = None |
| 852 | |
| 853 | |
| 854 | def aux_stream() -> torch.cuda.Stream | None: |
| 855 | """Ensures aux_stream is initialized only once.""" |
| 856 | global _aux_stream |
| 857 | |
| 858 | from vllm.platforms import current_platform |
| 859 | |
| 860 | if _aux_stream is None and current_platform.is_cuda_alike(): |
| 861 | _aux_stream = torch.cuda.Stream() |
| 862 | |
| 863 | return _aux_stream |
| 864 | |
| 865 | |
| 866 | def weak_ref_tensor(tensor: Any) -> Any: |
| 867 | """Create a weak reference to a tensor. |
| 868 | The new tensor will share the same data as the original tensor, |
| 869 | but will not keep the original tensor alive. |
| 870 | This ignores 0-size tensors as those don't allocate any memory. |
| 871 | """ |
| 872 | if isinstance(tensor, torch.Tensor) and tensor.numel() > 0: |
| 873 | return torch.ops._C.weak_ref_tensor(tensor) |
| 874 | else: |
| 875 | return tensor |
| 876 | |
| 877 | |
| 878 | def weak_ref_tensors( |
| 879 | tensors: torch.Tensor |
| 880 | | list[torch.Tensor] |
| 881 | | tuple[torch.Tensor] |
| 882 | | IntermediateTensors, |
| 883 | ) -> torch.Tensor | list[Any] | tuple[Any] | Any: |
| 884 | """Convenience function to create weak references to tensors, |
| 885 | for single tensor, list of tensors or tuple of tensors. |
| 886 | """ |
| 887 | if isinstance(tensors, torch.Tensor): |
| 888 | return weak_ref_tensor(tensors) |
| 889 | if isinstance(tensors, list): |
| 890 | return [weak_ref_tensor(t) for t in tensors] |
| 891 | if isinstance(tensors, tuple): |
| 892 | return tuple(weak_ref_tensor(t) for t in tensors) |
| 893 | |
| 894 | # For IntermediateTensors used in pipeline parallelism |
| 895 | from vllm.sequence import IntermediateTensors |
| 896 | |
| 897 | if isinstance(tensors, IntermediateTensors): |
| 898 | ret = IntermediateTensors( |
| 899 | {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()} |
| 900 | ) |
| 901 | return ret |
| 902 | raise ValueError("Invalid type for tensors") |
| 903 | |
| 904 | |
| 905 | def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor: |
| 906 | """Get an accelerator view of a CPU tensor via Unified Virtual Addressing.""" |
| 907 | from vllm.platforms import current_platform |
| 908 | |
| 909 | if current_platform.is_xpu(): |
| 910 | # Remove once the vllm-xpu-kernels fix for empty and non-pinned inputs |
| 911 | # (vllm-project/vllm-xpu-kernels#513) is in a released package. |
| 912 | if cpu_tensor.numel() == 0: |
| 913 | return torch.empty(cpu_tensor.shape, dtype=cpu_tensor.dtype, device="xpu") |
| 914 | if not cpu_tensor.is_pinned(): |
| 915 | pinned = torch.empty_strided( |
| 916 | cpu_tensor.size(), |
| 917 | cpu_tensor.stride(), |
| 918 | dtype=cpu_tensor.dtype, |
| 919 | device="cpu", |
| 920 | pin_memory=True, |
| 921 | ) |
| 922 | pinned.copy_(cpu_tensor) |
| 923 | return torch.ops._C.get_xpu_view_from_cpu_tensor(pinned) |
| 924 | return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor) |
| 925 | elif current_platform.is_cuda_alike(): |
| 926 | return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor) |
| 927 | else: |
| 928 | raise ValueError( |
| 929 | f"`get_accelerator_view_from_cpu_tensor` is currently " |
| 930 | f"not supported in: {current_platform.device_name}" |
| 931 | ) |
| 932 | |
| 933 | |
| 934 | # Helper function used in testing. |
| 935 | def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool: |
| 936 | return version.parse(torch_version) >= version.parse(target) |
| 937 | |
| 938 | |
| 939 | def is_torch_equal_or_newer(target: str) -> bool: |
| 940 | """Check if the installed torch version is >= the target version. |
| 941 | |
| 942 | Args: |
| 943 | target: a version string, like "2.6.0". |
| 944 | |
| 945 | Returns: |
| 946 | Whether the condition meets. |
| 947 | |
| 948 | """ |
| 949 | try: |
| 950 | return _is_torch_equal_or_newer(str(torch.__version__), target) |
| 951 | except Exception: |
| 952 | # Fallback to PKG-INFO to load the package info, needed by the doc gen. |
| 953 | return Version(importlib.metadata.version("torch")) >= Version(target) |
| 954 | |
| 955 | |
| 956 | def _is_torch_equal(target: str) -> bool: |
| 957 | assert target.count(".") == 2 |
| 958 | torch_version = str(torch.__version__) |
| 959 | torch_version = version.parse(torch_version) |
| 960 | # torch version is like "2.6.0.dev20240101" or "2.6.0.dev20240101+cpu" |
| 961 | # or "2.6.0+cu128" but never "2.6.0.1" |
| 962 | return ( |
| 963 | torch_version >= version.parse(target) |
| 964 | and version.parse(target + ".1") > torch_version |
| 965 | ) |
| 966 | |
| 967 | |
| 968 | def is_torch_equal(target: str) -> bool: |
| 969 | """Check if the installed torch version is == the target version. |
| 970 | |
| 971 | Args: |
| 972 | target: a version string, like "2.6.0". |
| 973 | |
| 974 | Returns: |
| 975 | Whether the condition meets. |
| 976 | |
| 977 | """ |
| 978 | try: |
| 979 | return _is_torch_equal(target) |
| 980 | except Exception: |
| 981 | return Version(importlib.metadata.version("torch")) == Version(target) |
| 982 | |
| 983 | |
| 984 | HAS_OPAQUE_TYPE = is_torch_equal_or_newer("2.11.0.dev") |
| 985 | |
| 986 | # Allow toggling LayerName usage via environment variable. |
| 987 | # Defaults to True on torch >= 2.11, False otherwise. |
| 988 | # Set VLLM_USE_LAYERNAME=0 to disable even on torch >= 2.11. |
| 989 | _USE_LAYERNAME = HAS_OPAQUE_TYPE and envs.VLLM_USE_LAYERNAME |
| 990 | |
| 991 | if HAS_OPAQUE_TYPE: |
| 992 | from torch._opaque_base import OpaqueBase |
| 993 | else: |
| 994 | OpaqueBase = object # type: ignore[misc, assignment] |
| 995 | |
| 996 | |
| 997 | class LayerName(OpaqueBase): # type: ignore[misc] |
| 998 | """Wraps a module name string for use as a torch opaque type. |
| 999 | |
| 1000 | When torch >= 2.11, this is registered as a hoisted value-type opaque |
| 1001 | object so that torch.compile lifts it as a graph input instead of baking |
| 1002 | it as a constant. This avoids per-layer recompilation for custom ops |
| 1003 | that accept layer name strings (attention, MOE, KV cache, etc.). |
| 1004 | """ |
| 1005 | |
| 1006 | def __init__(self, value: str): |
| 1007 | self.value = value |
| 1008 | |
| 1009 | def __eq__(self, other): |
| 1010 | return isinstance(other, LayerName) and self.value == other.value |
| 1011 | |
| 1012 | def __hash__(self): |
| 1013 | return hash(self.value) |
| 1014 | |
| 1015 | def __fx_repr__(self): |
| 1016 | return (f"LayerName({self.value!r})", {"LayerName": LayerName}) |
| 1017 | |
| 1018 | |
| 1019 | if HAS_OPAQUE_TYPE: |
| 1020 | from torch._library.opaque_object import register_opaque_type |
| 1021 | |
| 1022 | register_opaque_type(LayerName, typ="value", hoist=True) |
| 1023 | |
| 1024 | # On torch >= 2.11 (with VLLM_USE_LAYERNAME enabled), custom op |
| 1025 | # layer_name parameters use LayerName; otherwise they remain plain str. |
| 1026 | if TYPE_CHECKING: |
| 1027 | from typing import TypeAlias |
| 1028 | |
| 1029 | LayerNameType: TypeAlias = str | LayerName |
| 1030 | else: |
| 1031 | LayerNameType = LayerName if _USE_LAYERNAME else str |
| 1032 | |
| 1033 | |
| 1034 | def _resolve_layer_name(layer_name: str | LayerName) -> str: |
| 1035 | """Unwrap a LayerName to str, or return str unchanged.""" |
| 1036 | return layer_name.value if isinstance(layer_name, LayerName) else layer_name |
| 1037 | |
| 1038 | |
| 1039 | def _encode_layer_name(layer_name: str) -> str | LayerName: |
| 1040 | """Wrap a str layer name as LayerName when enabled.""" |
| 1041 | return LayerName(layer_name) if _USE_LAYERNAME else layer_name |
| 1042 | |
| 1043 | |
| 1044 | # Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform |
| 1045 | def supports_xpu_graph() -> bool: |
| 1046 | return is_torch_equal_or_newer("2.11.0.dev") |
| 1047 | |
| 1048 | |
| 1049 | # create a library to hold the custom op |
| 1050 | vllm_lib = Library("vllm", "FRAGMENT") # noqa |
| 1051 | |
| 1052 | |
| 1053 | def direct_register_custom_op( |
| 1054 | op_name: str, |
| 1055 | op_func: Callable, |
| 1056 | mutates_args: list[str] | None = None, |
| 1057 | fake_impl: Callable | None = None, |
| 1058 | target_lib: Library | None = None, |
| 1059 | dispatch_key: str | None = None, |
| 1060 | tags: tuple[torch.Tag, ...] = (), |
| 1061 | ): |
| 1062 | """`torch.library.custom_op` can have significant overhead because it |
| 1063 | needs to consider complicated dispatching logic. This function |
| 1064 | directly registers a custom op and dispatches it to the CUDA backend. |
| 1065 | See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5 |
| 1066 | for more details. |
| 1067 | |
| 1068 | By default, the custom op is registered to the vLLM library. If you |
| 1069 | want to register it to a different library, you can pass the library |
| 1070 | object to the `target_lib` argument. |
| 1071 | |
| 1072 | IMPORTANT: the lifetime of the operator is tied to the lifetime of the |
| 1073 | library object. If you want to bind the operator to a different library, |
| 1074 | make sure the library object is alive when the operator is used. |
| 1075 | """ |
| 1076 | if mutates_args is None: |
| 1077 | mutates_args = [] |
| 1078 | |
| 1079 | if dispatch_key is None: |
| 1080 | from vllm.platforms import current_platform |
| 1081 | |
| 1082 | dispatch_key = current_platform.dispatch_key |
| 1083 | |
| 1084 | schema_str = infer_schema(op_func, mutates_args=mutates_args) |
| 1085 | |
| 1086 | my_lib = target_lib or vllm_lib |
| 1087 | my_lib.define(op_name + schema_str, tags=tags) |
| 1088 | my_lib.impl(op_name, op_func, dispatch_key=dispatch_key) |
| 1089 | if fake_impl is not None: |
| 1090 | my_lib._register_fake(op_name, fake_impl) |