A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Chore] Clean up pytorch helper functions in `vllm.utils` (#26908)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: isotr0py <2037008807@qq.com>
Isotr0py committed
Oct 19, 2025 at 00:48 UTC
6ac5e06f7c5d4658c9fb119826a92d9910730fb4
119 files changed
+772
-714
benchmarks/kernels/bench_per_token_quant_fp8.py
+2
-1
@@ -10,7 +10,8 @@ import torch
10
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
11
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
12
from vllm.triton_utils import triton
13
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
13
+from vllm.utils import FlexibleArgumentParser
14
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
15
16
17
def with_triton_mode(fn):
benchmarks/kernels/benchmark_activation.py
+2
-1
@@ -10,7 +10,8 @@ import vllm.model_executor.layers.activation # noqa F401
10
from vllm.model_executor.custom_op import CustomOp
11
from vllm.platforms import current_platform
12
from vllm.triton_utils import triton
13
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
13
+from vllm.utils import FlexibleArgumentParser
14
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
15
16
batch_size_range = [1, 16, 32, 64, 128]
17
seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
benchmarks/kernels/benchmark_layernorm.py
+2
-1
@@ -7,7 +7,8 @@ import torch
7
8
from vllm.model_executor.layers.layernorm import RMSNorm
9
from vllm.platforms import current_platform
10
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
10
+from vllm.utils import FlexibleArgumentParser
11
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
12
13
14
@torch.inference_mode()
benchmarks/kernels/benchmark_paged_attention.py
+2
-2
@@ -9,9 +9,9 @@ import torch
9
from vllm import _custom_ops as ops
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
-from vllm.utils import (
12
+from vllm.utils import FlexibleArgumentParser
13
+from vllm.utils.torch_utils import (
14
STR_DTYPE_TO_TORCH_DTYPE,
14
- FlexibleArgumentParser,
15
create_kv_caches_with_random,
16
)
17
benchmarks/kernels/benchmark_quant.py
+2
-1
@@ -7,7 +7,8 @@ import torch
7
8
from vllm import _custom_ops as ops
9
from vllm.platforms import current_platform
10
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
10
+from vllm.utils import FlexibleArgumentParser
11
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
12
13
14
@torch.inference_mode()
benchmarks/kernels/benchmark_reshape_and_cache.py
+2
-2
@@ -9,9 +9,9 @@ from tabulate import tabulate
9
from vllm import _custom_ops as ops
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
-from vllm.utils import (
12
+from vllm.utils import FlexibleArgumentParser
13
+from vllm.utils.torch_utils import (
14
STR_DTYPE_TO_TORCH_DTYPE,
14
- FlexibleArgumentParser,
15
create_kv_caches_with_random,
16
)
17
benchmarks/kernels/benchmark_reshape_and_cache_flash.py
+2
-2
@@ -12,9 +12,9 @@ from vllm.attention.ops.triton_reshape_and_cache_flash import (
12
)
13
from vllm.logger import init_logger
14
from vllm.platforms import current_platform
15
-from vllm.utils import (
15
+from vllm.utils import FlexibleArgumentParser
16
+from vllm.utils.torch_utils import (
17
STR_DTYPE_TO_TORCH_DTYPE,
17
- FlexibleArgumentParser,
18
create_kv_caches_with_random_flash,
19
)
20
tests/compile/piecewise/test_full_cudagraph.py
+1
-1
@@ -11,7 +11,7 @@ from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
11
from vllm import LLM, SamplingParams
12
from vllm.config import CompilationConfig
13
from vllm.platforms import current_platform
14
-from vllm.utils import is_torch_equal_or_newer
14
+from vllm.utils.torch_utils import is_torch_equal_or_newer
15
16
17
@contextlib.contextmanager
tests/compile/piecewise/test_multiple_graphs.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.config import (
20
set_current_vllm_config,
21
)
22
from vllm.forward_context import BatchDescriptor, set_forward_context
23
-from vllm.utils import is_torch_equal_or_newer
23
+from vllm.utils.torch_utils import is_torch_equal_or_newer
24
25
# This import automatically registers `torch.ops.silly.attention`
26
from .. import silly_attention # noqa: F401
tests/compile/piecewise/test_simple.py
+1
-1
@@ -19,7 +19,7 @@ from vllm.config import (
19
set_current_vllm_config,
20
)
21
from vllm.forward_context import BatchDescriptor, set_forward_context
22
-from vllm.utils import is_torch_equal_or_newer
22
+from vllm.utils.torch_utils import is_torch_equal_or_newer
23
24
# This import automatically registers `torch.ops.silly.attention`
25
from ..silly_attention import get_global_counter, reset_global_counter
tests/compile/piecewise/test_toy_llama.py
+1
-1
@@ -27,7 +27,7 @@ from vllm.config import (
27
set_current_vllm_config,
28
)
29
from vllm.forward_context import BatchDescriptor, set_forward_context
30
-from vllm.utils import is_torch_equal_or_newer
30
+from vllm.utils.torch_utils import is_torch_equal_or_newer
31
32
# This import automatically registers `torch.ops.silly.attention`
33
from .. import silly_attention # noqa: F401
tests/compile/silly_attention.py
+1
-1
@@ -8,7 +8,7 @@ Centralizes custom operation definitions to avoid duplicate registrations.
8
import torch
9
from torch.library import Library
10
11
-from vllm.utils import direct_register_custom_op
11
+from vllm.utils.torch_utils import direct_register_custom_op
12
13
# Shared library for all compilation test operations
14
# Using "silly" namespace to match existing test expectations
tests/compile/test_aot_compile.py
+1
-1
@@ -15,7 +15,7 @@ from vllm.config import (
15
set_current_vllm_config,
16
)
17
from vllm.forward_context import set_forward_context
18
-from vllm.utils import is_torch_equal_or_newer
18
+from vllm.utils.torch_utils import is_torch_equal_or_newer
19
20
21
def reference_fn(x: torch.Tensor):
tests/compile/test_basic_correctness.py
+1
-1
@@ -5,7 +5,7 @@ import dataclasses
5
import pytest
6
7
from vllm.config import CompilationMode
8
-from vllm.utils import cuda_device_count_stateless
8
+from vllm.utils.torch_utils import cuda_device_count_stateless
9
10
from ..utils import compare_all_settings
11
tests/compile/test_config.py
+1
-1
@@ -8,7 +8,7 @@ from vllm.compilation.counter import compilation_counter
8
from vllm.compilation.fix_functionalization import FixFunctionalizationPass
9
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
10
from vllm.config.compilation import CompilationMode
11
-from vllm.utils import _is_torch_equal_or_newer, is_torch_equal_or_newer
11
+from vllm.utils.torch_utils import _is_torch_equal_or_newer, is_torch_equal_or_newer
12
13
14
def test_version():
tests/compile/test_decorator.py
+1
-1
@@ -15,7 +15,7 @@ from vllm.config import (
15
set_current_vllm_config,
16
)
17
from vllm.forward_context import BatchDescriptor, set_forward_context
18
-from vllm.utils import is_torch_equal_or_newer
18
+from vllm.utils.torch_utils import is_torch_equal_or_newer
19
20
# This import automatically registers `torch.ops.silly.attention`
21
from . import silly_attention # noqa: F401
tests/compile/test_full_graph.py
+1
-1
@@ -12,7 +12,7 @@ from tests.quantization.utils import is_quant_method_supported
12
from vllm import LLM, SamplingParams
13
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
14
from vllm.platforms import current_platform
15
-from vllm.utils import is_torch_equal_or_newer
15
+from vllm.utils.torch_utils import is_torch_equal_or_newer
16
17
from ..utils import create_new_process_for_each_test
18
tests/compile/test_fusions_e2e.py
+1
-1
@@ -15,8 +15,8 @@ from tests.v1.attention.utils import _Backend
15
from vllm import LLM, SamplingParams
16
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig
17
from vllm.platforms import current_platform
18
-from vllm.utils import is_torch_equal_or_newer
18
from vllm.utils.flashinfer import has_flashinfer
19
+from vllm.utils.torch_utils import is_torch_equal_or_newer
20
21
from ..utils import flat_product, multi_gpu_test
22
tests/conftest.py
+1
-1
@@ -60,8 +60,8 @@ from vllm.multimodal.utils import fetch_image
60
from vllm.outputs import RequestOutput
61
from vllm.sampling_params import BeamSearchParams
62
from vllm.transformers_utils.utils import maybe_model_redirect
63
-from vllm.utils import set_default_torch_num_threads
63
from vllm.utils.collections import is_list_of
64
+from vllm.utils.torch_utils import set_default_torch_num_threads
65
66
logger = init_logger(__name__)
67
tests/distributed/test_sequence_parallel.py
+1
-1
@@ -18,7 +18,7 @@ import pytest
18
from vllm.config.compilation import CompilationMode
19
from vllm.config.model import RunnerOption
20
from vllm.logger import init_logger
21
-from vllm.utils import is_torch_equal_or_newer
21
+from vllm.utils.torch_utils import is_torch_equal_or_newer
22
23
from ..models.registry import HF_EXAMPLE_MODELS
24
from ..utils import compare_two_settings, create_new_process_for_each_test
tests/distributed/test_utils.py
+1
-1
@@ -11,10 +11,10 @@ import vllm.envs as envs
11
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
12
from vllm.distributed.utils import StatelessProcessGroup
13
from vllm.utils import (
14
- cuda_device_count_stateless,
14
get_open_port,
15
update_environment_variables,
16
)
17
+from vllm.utils.torch_utils import cuda_device_count_stateless
18
19
from ..utils import multi_gpu_test
20
tests/kernels/attention/conftest.py
+4
-1
@@ -3,7 +3,10 @@
3
4
import pytest
5
6
-from vllm.utils import create_kv_caches_with_random, create_kv_caches_with_random_flash
6
+from vllm.utils.torch_utils import (
7
+ create_kv_caches_with_random,
8
+ create_kv_caches_with_random_flash,
9
+)
10
11
12
@pytest.fixture()
tests/kernels/attention/test_prefix_prefill.py
+1
-1
@@ -15,7 +15,7 @@ from tests.kernels.utils import make_alibi_bias
15
from vllm.attention.ops.chunked_prefill_paged_decode import chunked_prefill_paged_decode
16
from vllm.attention.ops.prefix_prefill import context_attention_fwd
17
from vllm.platforms import current_platform
18
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
18
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
19
20
NUM_HEADS = [64]
21
NUM_QUERIES_PER_KV = [1, 64]
tests/kernels/core/test_uva.py
+2
-1
@@ -3,7 +3,8 @@
3
import pytest
4
import torch
5
6
-from vllm.utils import get_cuda_view_from_cpu_tensor, is_uva_available
6
+from vllm.utils import is_uva_available
7
+from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
8
9
CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
10
tests/kernels/moe/test_modular_kernel_combinations.py
+2
-1
@@ -13,8 +13,9 @@ import torch
13
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14
from vllm.config import VllmConfig, set_current_vllm_config
15
from vllm.platforms import current_platform
16
-from vllm.utils import cuda_device_count_stateless, has_deep_ep, has_deep_gemm, has_pplx
16
+from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
17
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
18
+from vllm.utils.torch_utils import cuda_device_count_stateless
19
20
from .modular_kernel_tools.common import (
21
Config,
tests/kernels/utils.py
+1
-1
@@ -22,8 +22,8 @@ from vllm.utils import (
22
STR_BACKEND_ENV_VAR,
23
STR_FLASH_ATTN_VAL,
24
STR_XFORMERS_ATTN_VAL,
25
- make_tensor_with_pad,
25
)
26
+from vllm.utils.torch_utils import make_tensor_with_pad
27
28
# For now, disable "test_aot_dispatch_dynamic" since there are some
29
# bugs related to this test in PyTorch 2.4.
tests/models/multimodal/pooling/test_intern_vit.py
+1
-1
@@ -7,7 +7,7 @@ from huggingface_hub import snapshot_download
7
from transformers import AutoConfig, AutoModel, CLIPImageProcessor
8
9
from vllm.distributed import cleanup_dist_env_and_memory
10
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
10
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
11
12
from ....conftest import ImageTestAssets
13
tests/models/multimodal/pooling/test_radio.py
+1
-1
@@ -9,7 +9,7 @@ from transformers import AutoConfig, AutoModel, CLIPImageProcessor
9
from vllm.distributed import cleanup_dist_env_and_memory
10
from vllm.model_executor.models.radio import RadioModel
11
from vllm.transformers_utils.configs.radio import RadioConfig
12
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
12
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
13
14
from ....conftest import ImageTestAssets
15
tests/models/multimodal/processing/test_tensor_schema.py
+1
-1
@@ -26,7 +26,6 @@ from vllm.distributed import (
26
init_distributed_environment,
27
initialize_model_parallel,
28
)
29
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
29
from vllm.model_executor.models.interfaces import (
30
SupportsMultiModal,
31
supports_multimodal,
@@ -36,6 +35,7 @@ from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingC
35
from vllm.multimodal.utils import group_mm_kwargs_by_modality
36
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
37
from vllm.utils.collections import is_list_of
38
+from vllm.utils.torch_utils import set_default_torch_dtype
39
40
from ...registry import _MULTIMODAL_EXAMPLE_MODELS, HF_EXAMPLE_MODELS
41
from ...utils import dummy_hf_overrides
tests/utils.py
+1
-1
@@ -46,10 +46,10 @@ from vllm.platforms import current_platform
46
from vllm.transformers_utils.tokenizer import get_tokenizer
47
from vllm.utils import (
48
FlexibleArgumentParser,
49
- cuda_device_count_stateless,
49
get_open_port,
50
)
51
from vllm.utils.mem_constants import GB_bytes
52
+from vllm.utils.torch_utils import cuda_device_count_stateless
53
54
if current_platform.is_rocm():
55
from amdsmi import (
tests/utils_/test_utils.py
+6
-4
@@ -24,11 +24,8 @@ from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens
24
from vllm.utils import (
25
FlexibleArgumentParser,
26
bind_kv_cache,
27
- common_broadcastable_dtype,
28
- current_stream,
27
get_open_port,
28
get_tcp_uri,
31
- is_lossless_cast,
29
join_host_port,
30
make_zmq_path,
31
make_zmq_socket,
@@ -37,6 +34,11 @@ from vllm.utils import (
34
split_zmq_path,
35
unique_filepath,
36
)
37
+from vllm.utils.torch_utils import (
38
+ common_broadcastable_dtype,
39
+ current_stream,
40
+ is_lossless_cast,
41
+)
42
43
from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
44
from ..utils import create_new_process_for_each_test, flat_product
@@ -408,7 +410,7 @@ def test_bind_kv_cache_non_attention():
410
411
412
def test_bind_kv_cache_pp():
411
- with patch("vllm.utils.cuda_device_count_stateless", lambda: 2):
413
+ with patch("vllm.utils.torch_utils.cuda_device_count_stateless", lambda: 2):
414
# this test runs with 1 GPU, but we simulate 2 GPUs
415
cfg = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=2))
416
with set_current_vllm_config(cfg):
tests/v1/attention/test_attention_backends.py
+2
-1
@@ -18,7 +18,8 @@ from tests.v1.attention.utils import (
18
from vllm.attention.backends.registry import _Backend
19
from vllm.config import ModelConfig
20
from vllm.platforms import current_platform
21
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv, is_torch_equal_or_newer
21
+from vllm.utils import cdiv
22
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, is_torch_equal_or_newer
23
from vllm.v1.attention.backends.utils import (
24
CommonAttentionMetadata,
25
set_kv_cache_layout,
tests/v1/attention/test_mla_backends.py
+2
-1
@@ -22,7 +22,8 @@ from vllm import _custom_ops as ops
22
from vllm.attention.backends.registry import _Backend
23
from vllm.attention.ops.flashmla import is_flashmla_dense_supported
24
from vllm.config.vllm import set_current_vllm_config
25
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
25
+from vllm.utils import cdiv
26
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
27
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
28
from vllm.v1.kv_cache_interface import FullAttentionSpec
29
tests/v1/engine/test_async_llm.py
+1
-1
@@ -15,7 +15,7 @@ from vllm.inputs import PromptType
15
from vllm.outputs import RequestOutput
16
from vllm.platforms import current_platform
17
from vllm.sampling_params import RequestOutputKind
18
-from vllm.utils import set_default_torch_num_threads
18
+from vllm.utils.torch_utils import set_default_torch_num_threads
19
from vllm.v1.engine.async_llm import AsyncLLM
20
from vllm.v1.metrics.loggers import (
21
AggregatedLoggingStatLogger,
tests/v1/engine/test_engine_core.py
+1
-1
@@ -12,7 +12,7 @@ from transformers import AutoTokenizer
12
from vllm import SamplingParams
13
from vllm.engine.arg_utils import EngineArgs
14
from vllm.platforms import current_platform
15
-from vllm.utils import set_default_torch_num_threads
15
+from vllm.utils.torch_utils import set_default_torch_num_threads
16
from vllm.v1.engine import EngineCoreRequest
17
from vllm.v1.engine.core import EngineCore
18
from vllm.v1.executor.abstract import Executor, UniProcExecutor
tests/v1/engine/test_engine_core_client.py
+1
-1
@@ -21,7 +21,7 @@ from vllm.distributed.kv_events import BlockStored, KVEventBatch, ZmqEventPublis
21
from vllm.engine.arg_utils import EngineArgs
22
from vllm.platforms import current_platform
23
from vllm.usage.usage_lib import UsageContext
24
-from vllm.utils import set_default_torch_num_threads
24
+from vllm.utils.torch_utils import set_default_torch_num_threads
25
from vllm.v1.engine import EngineCoreRequest
26
from vllm.v1.engine.core import EngineCore
27
from vllm.v1.engine.core_client import AsyncMPClient, EngineCoreClient, SyncMPClient
tests/v1/sample/test_sampler.py
+2
-1
@@ -7,7 +7,8 @@ import torch
7
8
from tests.v1.sample.utils import create_allowed_token_ids
9
from vllm.platforms import current_platform
10
-from vllm.utils import is_pin_memory_available, make_tensor_with_pad
10
+from vllm.utils import is_pin_memory_available
11
+from vllm.utils.torch_utils import make_tensor_with_pad
12
from vllm.v1.sample.logits_processor import LogitsProcessors
13
from vllm.v1.sample.metadata import SamplingMetadata
14
from vllm.v1.sample.sampler import Sampler
tests/v1/sample/utils.py
+1
-1
@@ -9,7 +9,7 @@ import regex as re
9
import torch
10
11
from vllm import CompletionOutput
12
-from vllm.utils import make_tensor_with_pad
12
+from vllm.utils.torch_utils import make_tensor_with_pad
13
from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor
14
from vllm.v1.sample.metadata import SamplingMetadata
15
tests/v1/shutdown/test_delete.py
+1
-1
@@ -12,7 +12,7 @@ from tests.v1.shutdown.utils import (
12
from vllm import LLM, SamplingParams
13
from vllm.engine.arg_utils import AsyncEngineArgs
14
from vllm.sampling_params import RequestOutputKind
15
-from vllm.utils import cuda_device_count_stateless
15
+from vllm.utils.torch_utils import cuda_device_count_stateless
16
from vllm.v1.engine.async_llm import AsyncLLM
17
18
MODELS = ["hmellor/tiny-random-LlamaForCausalLM"]
tests/v1/shutdown/test_forward_error.py
+1
-1
@@ -14,7 +14,7 @@ from tests.v1.shutdown.utils import (
14
from vllm import LLM, AsyncEngineArgs, SamplingParams
15
from vllm.distributed import get_tensor_model_parallel_rank
16
from vllm.model_executor.models.llama import LlamaForCausalLM
17
-from vllm.utils import cuda_device_count_stateless
17
+from vllm.utils.torch_utils import cuda_device_count_stateless
18
from vllm.v1.engine.async_llm import AsyncLLM
19
from vllm.v1.engine.exceptions import EngineDeadError
20
tests/v1/shutdown/test_startup_error.py
+1
-1
@@ -13,7 +13,7 @@ from vllm import LLM
13
from vllm.distributed import get_tensor_model_parallel_rank
14
from vllm.engine.arg_utils import AsyncEngineArgs
15
from vllm.model_executor.models.llama import LlamaForCausalLM
16
-from vllm.utils import cuda_device_count_stateless
16
+from vllm.utils.torch_utils import cuda_device_count_stateless
17
from vllm.v1.engine.async_llm import AsyncLLM
18
19
MODELS = ["hmellor/tiny-random-LlamaForCausalLM"]
tests/v1/worker/test_gpu_input_batch.py
+2
-1
@@ -10,7 +10,8 @@ import torch
10
11
from vllm.platforms import current_platform
12
from vllm.sampling_params import SamplingParams
13
-from vllm.utils import is_pin_memory_available, make_tensor_with_pad
13
+from vllm.utils import is_pin_memory_available
14
+from vllm.utils.torch_utils import make_tensor_with_pad
15
from vllm.v1.pool.metadata import PoolingMetadata
16
from vllm.v1.sample.logits_processor import LogitsProcessors
17
from vllm.v1.sample.metadata import SamplingMetadata
vllm/attention/layer.py
+1
-1
@@ -35,7 +35,7 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
35
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
36
from vllm.model_executor.models.vision import get_vit_attn_backend
37
from vllm.platforms import current_platform
38
-from vllm.utils import (
38
+from vllm.utils.torch_utils import (
39
direct_register_custom_op,
40
kv_cache_dtype_str_to_dtype,
41
)
vllm/attention/ops/rocm_aiter_mla.py
+1
-1
@@ -5,7 +5,7 @@
5
import torch
6
7
from vllm.platforms import current_platform
8
-from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
8
+from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
9
10
11
def get_aiter_mla_metadata(
vllm/compilation/backends.py
+1
-1
@@ -24,8 +24,8 @@ from vllm.compilation.partition_rules import (
24
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
25
from vllm.logger import init_logger
26
from vllm.platforms import current_platform
27
-from vllm.utils import is_torch_equal_or_newer
27
from vllm.utils.import_utils import resolve_obj_by_qualname
28
+from vllm.utils.torch_utils import is_torch_equal_or_newer
29
30
from .caching import VllmSerializableFunction
31
from .compiler_interface import (
vllm/compilation/collective_fusion.py
+1
-1
@@ -21,7 +21,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
kFp8StaticTensorSym,
22
)
23
from vllm.platforms import current_platform
24
-from vllm.utils import direct_register_custom_op
24
+from vllm.utils.torch_utils import direct_register_custom_op
25
26
from .inductor_pass import enable_fake_mode
27
from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
vllm/compilation/compiler_interface.py
+1
-1
@@ -16,7 +16,7 @@ import torch.fx as fx
16
import vllm.envs as envs
17
from vllm.compilation.counter import compilation_counter
18
from vllm.config import VllmConfig
19
-from vllm.utils import is_torch_equal_or_newer
19
+from vllm.utils.torch_utils import is_torch_equal_or_newer
20
21
22
class CompilerInterface:
vllm/compilation/cuda_graph.py
+1
-1
@@ -17,7 +17,7 @@ from vllm.distributed.device_communicators.pynccl_allocator import set_graph_poo
17
from vllm.forward_context import BatchDescriptor, get_forward_context
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
-from vllm.utils import weak_ref_tensors
20
+from vllm.utils.torch_utils import weak_ref_tensors
21
22
logger = init_logger(__name__)
23
vllm/compilation/decorators.py
+1
-1
@@ -21,8 +21,8 @@ from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
21
from vllm.config import CompilationMode, VllmConfig, set_current_vllm_config
22
from vllm.logger import init_logger
23
from vllm.sequence import IntermediateTensors
24
-from vllm.utils import supports_dynamo
24
from vllm.utils.import_utils import resolve_obj_by_qualname
25
+from vllm.utils.torch_utils import supports_dynamo
26
27
from .monitor import start_monitoring_torch_compile
28
vllm/compilation/inductor_pass.py
+1
-1
@@ -14,7 +14,7 @@ import torch
14
from torch import fx
15
from torch._subclasses.fake_tensor import FakeTensorMode, unset_fake_temporarily
16
17
-from vllm.utils import is_torch_equal_or_newer
17
+from vllm.utils.torch_utils import is_torch_equal_or_newer
18
19
if is_torch_equal_or_newer("2.6"):
20
from torch._inductor.custom_graph_pass import CustomGraphPass
vllm/config/compilation.py
+1
-1
@@ -16,8 +16,8 @@ from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
16
from vllm.config.utils import config
17
from vllm.logger import init_logger
18
from vllm.platforms import current_platform
19
-from vllm.utils import is_torch_equal_or_newer
19
from vllm.utils.import_utils import resolve_obj_by_qualname
20
+from vllm.utils.torch_utils import is_torch_equal_or_newer
21
22
if TYPE_CHECKING:
23
from vllm.config import VllmConfig
vllm/config/model.py
+2
-1
@@ -41,8 +41,9 @@ from vllm.transformers_utils.config import (
41
)
42
from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
43
from vllm.transformers_utils.utils import maybe_model_redirect
44
-from vllm.utils import LayerBlockType, common_broadcastable_dtype
44
+from vllm.utils import LayerBlockType
45
from vllm.utils.import_utils import LazyLoader
46
+from vllm.utils.torch_utils import common_broadcastable_dtype
47
48
if TYPE_CHECKING:
49
from transformers import PretrainedConfig
vllm/config/parallel.py
+2
-1
@@ -18,7 +18,8 @@ from vllm.model_executor.layers.batch_invariant import (
18
vllm_is_batch_invariant,
19
)
20
from vllm.platforms import current_platform
21
-from vllm.utils import cuda_device_count_stateless, get_open_ports_list
21
+from vllm.utils import get_open_ports_list
22
+from vllm.utils.torch_utils import cuda_device_count_stateless
23
24
if TYPE_CHECKING:
25
from ray.runtime_env import RuntimeEnv
vllm/distributed/device_communicators/all_reduce_utils.py
+2
-1
@@ -22,7 +22,8 @@ from vllm.logger import init_logger
22
from vllm.model_executor.layers.batch_invariant import (
23
vllm_is_batch_invariant,
24
)
25
-from vllm.utils import cuda_device_count_stateless, update_environment_variables
25
+from vllm.utils import update_environment_variables
26
+from vllm.utils.torch_utils import cuda_device_count_stateless
27
28
logger = init_logger(__name__)
29
vllm/distributed/device_communicators/custom_all_reduce.py
+1
-1
@@ -17,7 +17,7 @@ from vllm.distributed.device_communicators.all_reduce_utils import (
17
from vllm.distributed.parallel_state import in_the_same_node_as
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
-from vllm.utils import cuda_device_count_stateless
20
+from vllm.utils.torch_utils import cuda_device_count_stateless
21
22
try:
23
ops.meta_size()
vllm/distributed/device_communicators/pynccl.py
+2
-2
@@ -19,7 +19,7 @@ from vllm.distributed.device_communicators.pynccl_wrapper import (
19
)
20
from vllm.distributed.utils import StatelessProcessGroup
21
from vllm.logger import init_logger
22
-from vllm.utils import current_stream
22
+from vllm.utils.torch_utils import current_stream
23
24
logger = init_logger(__name__)
25
@@ -30,7 +30,7 @@ def register_nccl_symmetric_ops(pynccl_comm):
30
from vllm.distributed.device_communicators.pynccl_allocator import (
31
nccl_symm_mem_context,
32
)
33
- from vllm.utils import direct_register_custom_op
33
+ from vllm.utils.torch_utils import direct_register_custom_op
34
35
global _NCCL_SYMM_OPS_REGISTERED
36
if _NCCL_SYMM_OPS_REGISTERED:
vllm/distributed/device_communicators/quick_all_reduce.py
+1
-1
@@ -13,7 +13,7 @@ from vllm.config import get_current_vllm_config
13
from vllm.distributed.parallel_state import in_the_same_node_as
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
-from vllm.utils import cuda_device_count_stateless
16
+from vllm.utils.torch_utils import cuda_device_count_stateless
17
18
logger = init_logger(__name__)
19
vllm/distributed/device_communicators/ray_communicator.py
+1
-1
@@ -14,7 +14,7 @@ from vllm.distributed.device_communicators.base_device_communicator import (
14
)
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.logger import init_logger
17
-from vllm.utils import current_stream
17
+from vllm.utils.torch_utils import current_stream
18
19
logger = init_logger(__name__)
20
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+2
-1
@@ -25,7 +25,8 @@ from vllm.distributed.device_communicators.pynccl_wrapper import (
25
from vllm.distributed.kv_transfer.kv_connector.v1.p2p.tensor_memory_pool import ( # noqa: E501
26
TensorMemoryPool,
27
)
28
-from vllm.utils import current_stream, get_ip
28
+from vllm.utils import get_ip
29
+from vllm.utils.torch_utils import current_stream
30
31
logger = logging.getLogger(__name__)
32
vllm/distributed/parallel_state.py
+4
-2
@@ -50,11 +50,13 @@ from vllm.distributed.device_communicators.base_device_communicator import (
50
from vllm.distributed.utils import StatelessProcessGroup
51
from vllm.logger import init_logger
52
from vllm.utils import (
53
- direct_register_custom_op,
53
get_distributed_init_method,
55
- supports_custom_op,
54
)
55
from vllm.utils.import_utils import resolve_obj_by_qualname
56
+from vllm.utils.torch_utils import (
57
+ direct_register_custom_op,
58
+ supports_custom_op,
59
+)
60
61
62
@dataclass
vllm/distributed/utils.py
+2
-1
@@ -29,7 +29,8 @@ from torch.distributed.rendezvous import rendezvous
29
30
import vllm.envs as envs
31
from vllm.logger import init_logger
32
-from vllm.utils import get_tcp_uri, is_torch_equal_or_newer
32
+from vllm.utils import get_tcp_uri
33
+from vllm.utils.torch_utils import is_torch_equal_or_newer
34
35
logger = init_logger(__name__)
36
vllm/env_override.py
+1
-1
@@ -5,7 +5,7 @@ import os
5
import torch
6
7
from vllm.logger import init_logger
8
-from vllm.utils import is_torch_equal
8
+from vllm.utils.torch_utils import is_torch_equal
9
10
logger = init_logger(__name__)
11
vllm/envs.py
+1
-1
@@ -246,7 +246,7 @@ def maybe_convert_bool(value: str | None) -> bool | None:
246
247
248
def use_aot_compile() -> bool:
249
- from vllm.utils import is_torch_equal_or_newer
249
+ from vllm.utils.torch_utils import is_torch_equal_or_newer
250
251
default_value = "1" if is_torch_equal_or_newer("2.10.0.dev") else "0"
252
return os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1"
vllm/lora/ops/triton_ops/lora_expand_op.py
+1
-1
@@ -12,7 +12,7 @@ import torch
12
from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
13
from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr, get_lora_op_configs
14
from vllm.triton_utils import tl, triton
15
-from vllm.utils import direct_register_custom_op
15
+from vllm.utils.torch_utils import direct_register_custom_op
16
17
18
@triton.jit
vllm/lora/ops/triton_ops/lora_shrink_op.py
+1
-1
@@ -12,7 +12,7 @@ import torch
12
from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel
13
from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr, get_lora_op_configs
14
from vllm.triton_utils import tl, triton
15
-from vllm.utils import direct_register_custom_op
15
+from vllm.utils.torch_utils import direct_register_custom_op
16
17
18
@triton.jit
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+1
-1
@@ -10,7 +10,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
10
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
11
per_token_group_quant_fp8,
12
)
13
-from vllm.utils import direct_register_custom_op
13
+from vllm.utils.torch_utils import direct_register_custom_op
14
15
16
def flashinfer_fused_moe_blockscale_fp8(
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-1
@@ -52,8 +52,8 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Sc
52
from vllm.model_executor.utils import maybe_disable_graph_partition
53
from vllm.platforms import current_platform
54
from vllm.triton_utils import tl, triton
55
-from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
55
from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
56
+from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
57
58
from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled
59
vllm/model_executor/layers/fused_moe/layer.py
+2
-1
@@ -52,8 +52,9 @@ from vllm.model_executor.layers.quantization.base_config import (
52
from vllm.model_executor.utils import set_weight_attrs
53
from vllm.platforms import current_platform
54
from vllm.platforms.interface import CpuArchEnum
55
-from vllm.utils import cdiv, direct_register_custom_op, has_deep_ep, has_pplx, round_up
55
+from vllm.utils import cdiv, has_deep_ep, has_pplx, round_up
56
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
57
+from vllm.utils.torch_utils import direct_register_custom_op
58
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
59
60
if current_platform.is_cuda_alike():
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+1
-1
@@ -11,7 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
)
13
from vllm.platforms import current_platform
14
-from vllm.utils import direct_register_custom_op
14
+from vllm.utils.torch_utils import direct_register_custom_op
15
16
17
class QuantMethod(IntEnum):
vllm/model_executor/layers/fused_moe/utils.py
+2
-1
@@ -23,8 +23,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
23
mxfp8_e4m3_quantize,
24
)
25
from vllm.triton_utils import tl, triton
26
-from vllm.utils import cdiv, is_torch_equal_or_newer
26
+from vllm.utils import cdiv
27
from vllm.utils.flashinfer import flashinfer_fp4_quantize
28
+from vllm.utils.torch_utils import is_torch_equal_or_newer
29
30
31
@triton.jit
vllm/model_executor/layers/layernorm.py
+1
-1
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.batch_invariant import (
13
vllm_is_batch_invariant,
14
)
15
from vllm.platforms import current_platform
16
-from vllm.utils import direct_register_custom_op
16
+from vllm.utils.torch_utils import direct_register_custom_op
17
18
19
def is_rocm_aiter_rmsnorm_enabled() -> bool:
vllm/model_executor/layers/mamba/linear_attn.py
+1
-1
@@ -34,7 +34,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
34
MambaStateShapeCalculator,
35
)
36
from vllm.model_executor.layers.quantization import QuantizationConfig
37
-from vllm.utils import direct_register_custom_op
37
+from vllm.utils.torch_utils import direct_register_custom_op
38
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
39
40
if TYPE_CHECKING:
vllm/model_executor/layers/mamba/mamba_mixer.py
+1
-1
@@ -37,7 +37,7 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
37
selective_state_update,
38
)
39
from vllm.model_executor.utils import set_weight_attrs
40
-from vllm.utils import direct_register_custom_op
40
+from vllm.utils.torch_utils import direct_register_custom_op
41
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
42
43
vllm/model_executor/layers/mamba/mamba_mixer2.py
+1
-1
@@ -46,7 +46,7 @@ from vllm.model_executor.model_loader.weight_utils import (
46
sharded_weight_loader,
47
)
48
from vllm.model_executor.utils import set_weight_attrs
49
-from vllm.utils import direct_register_custom_op
49
+from vllm.utils.torch_utils import direct_register_custom_op
50
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
51
52
# Added by the IBM Team, 2024
vllm/model_executor/layers/mamba/mamba_utils.py
+4
-1
@@ -6,7 +6,10 @@ import torch
6
from vllm.config.cache import MambaDType
7
from vllm.config.model import ModelDType
8
from vllm.distributed import divide
9
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, get_kv_cache_torch_dtype
9
+from vllm.utils.torch_utils import (
10
+ STR_DTYPE_TO_TORCH_DTYPE,
11
+ get_kv_cache_torch_dtype,
12
+)
13
14
15
class MambaStateDtypeCalculator:
vllm/model_executor/layers/mamba/short_conv.py
+1
-1
@@ -27,7 +27,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
27
causal_conv1d_fn,
28
causal_conv1d_update,
29
)
30
-from vllm.utils import direct_register_custom_op
30
+from vllm.utils.torch_utils import direct_register_custom_op
31
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata
32
33
vllm/model_executor/layers/quantization/bitsandbytes.py
+1
-1
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization import (
23
QuantizationMethods,
24
)
25
from vllm.platforms import current_platform
26
-from vllm.utils import direct_register_custom_op
26
+from vllm.utils.torch_utils import direct_register_custom_op
27
28
29
class BitsAndBytesConfig(QuantizationConfig):
vllm/model_executor/layers/quantization/fp_quant.py
+1
-1
@@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizationConf
24
from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
25
from vllm.model_executor.utils import set_weight_attrs
26
from vllm.platforms import current_platform
27
-from vllm.utils import direct_register_custom_op
27
+from vllm.utils.torch_utils import direct_register_custom_op
28
29
30
class FPQuantConfig(QuantizationConfig):
vllm/model_executor/layers/quantization/gguf.py
+1
-1
@@ -28,7 +28,7 @@ from vllm.model_executor.layers.quantization.base_config import (
28
)
29
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
30
from vllm.model_executor.utils import set_weight_attrs
31
-from vllm.utils import direct_register_custom_op
31
+from vllm.utils.torch_utils import direct_register_custom_op
32
33
logger = init_logger(__name__)
34
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
+1
-1
@@ -7,7 +7,7 @@ import torch
7
import vllm.envs as envs
8
from vllm import _custom_ops as ops
9
from vllm.platforms import current_platform
10
-from vllm.utils import direct_register_custom_op
10
+from vllm.utils.torch_utils import direct_register_custom_op
11
12
from .cutlass import CutlassScaledMMLinearKernel
13
from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
vllm/model_executor/layers/quantization/mxfp4.py
+1
-1
@@ -49,10 +49,10 @@ from vllm.platforms import current_platform
49
from vllm.scalar_type import scalar_types
50
from vllm.utils import (
51
has_triton_kernels,
52
- is_torch_equal_or_newer,
52
round_up,
53
)
54
from vllm.utils.flashinfer import has_flashinfer
55
+from vllm.utils.torch_utils import is_torch_equal_or_newer
56
57
logger = init_logger(__name__)
58
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+1
-1
@@ -45,7 +45,7 @@ try:
45
from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
46
from aiter.ops.triton.quant import dynamic_mxfp4_quant
47
48
- from vllm.utils import direct_register_custom_op
48
+ from vllm.utils.torch_utils import direct_register_custom_op
49
50
if is_rocm_aiter_fp4_asm_gemm_enabled():
51
from aiter import gemm_a4w4, per_1x32_f4_quant_hip
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
-1
@@ -28,13 +28,13 @@ from vllm.model_executor.parameter import (
28
)
29
from vllm.platforms import current_platform
30
from vllm.triton_utils import tl, triton
31
-from vllm.utils import direct_register_custom_op
31
from vllm.utils.deep_gemm import (
32
fp8_gemm_nt,
33
is_deep_gemm_e8m0_used,
34
is_deep_gemm_supported,
35
should_use_deepgemm_for_fp8_linear,
36
)
37
+from vllm.utils.torch_utils import direct_register_custom_op
38
39
logger = init_logger(__name__)
40
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+1
-1
@@ -7,7 +7,7 @@ import torch
7
8
from vllm.logger import init_logger
9
from vllm.platforms import current_platform
10
-from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
10
+from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
11
12
logger = init_logger(__name__)
13
vllm/model_executor/layers/quantization/utils/mxfp6_utils.py
+1
-1
@@ -3,7 +3,7 @@
3
import torch
4
5
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE
6
-from vllm.utils import direct_register_custom_op
6
+from vllm.utils.torch_utils import direct_register_custom_op
7
8
9
def _quant_dequant_mxfp6(
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+1
-1
@@ -12,8 +12,8 @@ from vllm.config import CompilationMode, get_current_vllm_config
12
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
13
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
14
from vllm.platforms import current_platform
15
-from vllm.utils import direct_register_custom_op
15
from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm, has_flashinfer
16
+from vllm.utils.torch_utils import direct_register_custom_op
17
18
# Input scaling factors are no longer optional in _scaled_mm starting
19
# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
vllm/model_executor/layers/rotary_embedding/common.py
+1
-1
@@ -10,7 +10,7 @@ import torch
10
11
from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
-from vllm.utils import direct_register_custom_op
13
+from vllm.utils.torch_utils import direct_register_custom_op
14
15
if current_platform.is_cuda():
16
from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
vllm/model_executor/layers/rotary_embedding/rocm_aiter_rope_ops.py
+1
-1
@@ -5,7 +5,7 @@ import torch
5
6
import vllm.envs as envs
7
from vllm.platforms import current_platform
8
-from vllm.utils import direct_register_custom_op
8
+from vllm.utils.torch_utils import direct_register_custom_op
9
10
11
def is_rocm_triton_rotary_embedding_enabled() -> bool:
vllm/model_executor/layers/utils.py
+1
-1
@@ -9,7 +9,7 @@ import torch
9
from vllm import _custom_ops as ops
10
from vllm import envs
11
from vllm.platforms import CpuArchEnum, current_platform
12
-from vllm.utils import direct_register_custom_op
12
+from vllm.utils.torch_utils import direct_register_custom_op
13
14
15
def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
vllm/model_executor/model_loader/base_loader.py
+1
-1
@@ -11,8 +11,8 @@ from vllm.logger import init_logger
11
from vllm.model_executor.model_loader.utils import (
12
initialize_model,
13
process_weights_after_loading,
14
- set_default_torch_dtype,
14
)
15
+from vllm.utils.torch_utils import set_default_torch_dtype
16
17
logger = init_logger(__name__)
18
vllm/model_executor/model_loader/bitsandbytes_loader.py
+2
-1
@@ -32,7 +32,7 @@ from vllm.model_executor.layers.linear import (
32
RowParallelLinear,
33
)
34
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
35
-from vllm.model_executor.model_loader.utils import ParamMapping, set_default_torch_dtype
35
+from vllm.model_executor.model_loader.utils import ParamMapping
36
from vllm.model_executor.model_loader.weight_utils import (
37
download_safetensors_index_file_from_hf,
38
download_weights_from_hf,
@@ -48,6 +48,7 @@ from vllm.model_executor.utils import (
48
set_weight_attrs,
49
)
50
from vllm.platforms import current_platform
51
+from vllm.utils.torch_utils import set_default_torch_dtype
52
53
logger = init_logger(__name__)
54
vllm/model_executor/model_loader/gguf_loader.py
+1
-1
@@ -15,13 +15,13 @@ from vllm.model_executor.model_loader.base_loader import BaseModelLoader
15
from vllm.model_executor.model_loader.utils import (
16
initialize_model,
17
process_weights_after_loading,
18
- set_default_torch_dtype,
18
)
19
from vllm.model_executor.model_loader.weight_utils import (
20
get_gguf_extra_tensor_names,
21
get_gguf_weight_type_map,
22
gguf_quant_weights_iterator,
23
)
24
+from vllm.utils.torch_utils import set_default_torch_dtype
25
26
27
class GGUFModelLoader(BaseModelLoader):
vllm/model_executor/model_loader/tensorizer_loader.py
+1
-1
@@ -22,8 +22,8 @@ from vllm.model_executor.model_loader.tensorizer import (
22
from vllm.model_executor.model_loader.utils import (
23
get_model_architecture,
24
initialize_model,
25
- set_default_torch_dtype,
25
)
26
+from vllm.utils.torch_utils import set_default_torch_dtype
27
28
logger = init_logger(__name__)
29
vllm/model_executor/model_loader/tpu.py
+1
-1
@@ -14,8 +14,8 @@ from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
14
from vllm.model_executor.model_loader.utils import (
15
initialize_model,
16
process_weights_after_loading,
17
- set_default_torch_dtype,
17
)
18
+from vllm.utils.torch_utils import set_default_torch_dtype
19
20
logger = init_logger(__name__)
21
vllm/model_executor/model_loader/utils.py
-10
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Utilities for selecting and loading models."""
4
5
-import contextlib
5
import inspect
6
import warnings
7
from contextlib import contextmanager
@@ -32,15 +31,6 @@ from vllm.utils import is_pin_memory_available
31
logger = init_logger(__name__)
32
33
35
-@contextlib.contextmanager
36
-def set_default_torch_dtype(dtype: torch.dtype):
37
- """Sets the default torch dtype to the given dtype."""
38
- old_dtype = torch.get_default_dtype()
39
- torch.set_default_dtype(dtype)
40
- yield
41
- torch.set_default_dtype(old_dtype)
42
-
43
-
34
def initialize_model(
35
vllm_config: VllmConfig,
36
*,
vllm/model_executor/models/config.py
+2
-1
@@ -6,7 +6,8 @@ from typing import TYPE_CHECKING
6
import vllm.envs as envs
7
from vllm.logger import init_logger
8
from vllm.model_executor.models import ModelRegistry
9
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv, round_up
9
+from vllm.utils import cdiv, round_up
10
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
11
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
12
13
if TYPE_CHECKING:
vllm/model_executor/models/deepseek_v2.py
+1
-1
@@ -79,8 +79,8 @@ from vllm.model_executor.model_loader.weight_utils import (
79
from vllm.model_executor.models.utils import sequence_parallel_chunk
80
from vllm.platforms import current_platform
81
from vllm.sequence import IntermediateTensors
82
-from vllm.utils import direct_register_custom_op
82
from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
83
+from vllm.utils.torch_utils import direct_register_custom_op
84
from vllm.v1.attention.backends.mla.indexer import (
85
DeepseekV32IndexerBackend,
86
DeepseekV32IndexerMetadata,
vllm/model_executor/models/deepseek_vl2.py
+1
-1
@@ -18,7 +18,6 @@ from vllm.config import VllmConfig
18
from vllm.config.multimodal import BaseDummyOptions
19
from vllm.distributed import get_tensor_model_parallel_world_size
20
from vllm.model_executor.layers.quantization import QuantizationConfig
21
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
21
from vllm.model_executor.models.transformers.utils import replace_linear_class
22
from vllm.multimodal import MULTIMODAL_REGISTRY
23
from vllm.multimodal.inputs import (
@@ -51,6 +50,7 @@ from vllm.transformers_utils.processors.deepseek_vl2 import DeepseekVLV2Processo
50
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
51
from vllm.utils.collections import is_list_of
52
from vllm.utils.tensor_schema import TensorSchema, TensorShape
53
+from vllm.utils.torch_utils import set_default_torch_dtype
54
55
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
56
from .utils import (
vllm/model_executor/models/internvl.py
+1
-1
@@ -51,8 +51,8 @@ from vllm.multimodal.processing import (
51
from vllm.multimodal.profiling import BaseDummyInputsBuilder
52
from vllm.sequence import IntermediateTensors
53
from vllm.transformers_utils.tokenizer import AnyTokenizer
54
-from vllm.utils import set_default_torch_num_threads
54
from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
+from vllm.utils.torch_utils import set_default_torch_num_threads
56
57
from .interfaces import (
58
MultiModalEmbeddings,
vllm/model_executor/models/minicpmv.py
+1
-1
@@ -49,7 +49,6 @@ from vllm.model_executor.layers.resampler import (
49
Resampler2,
50
get_2d_sincos_pos_embed,
51
)
52
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
52
from vllm.model_executor.models.llama import LlamaForCausalLM
53
from vllm.model_executor.models.minicpm import MiniCPMForCausalLM
54
from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -88,6 +87,7 @@ from vllm.platforms import current_platform
87
from vllm.sequence import IntermediateTensors
88
from vllm.utils.collections import flatten_2d_lists
89
from vllm.utils.tensor_schema import TensorSchema, TensorShape
90
+from vllm.utils.torch_utils import set_default_torch_dtype
91
92
from .idefics2_vision_model import Idefics2VisionTransformer
93
from .interfaces import (
vllm/model_executor/models/plamo2.py
+1
-1
@@ -64,7 +64,7 @@ from vllm.model_executor.models.utils import (
64
)
65
from vllm.model_executor.utils import set_weight_attrs
66
from vllm.sequence import IntermediateTensors
67
-from vllm.utils import direct_register_custom_op
67
+from vllm.utils.torch_utils import direct_register_custom_op
68
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
69
70
vllm/model_executor/models/qwen3_next.py
+1
-1
@@ -71,7 +71,7 @@ from vllm.platforms import current_platform
71
from vllm.sequence import IntermediateTensors
72
from vllm.transformers_utils.configs import Qwen3NextConfig
73
from vllm.triton_utils import tl, triton
74
-from vllm.utils import direct_register_custom_op
74
+from vllm.utils.torch_utils import direct_register_custom_op
75
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
76
77
from .interfaces import (
vllm/model_executor/models/transformers/moe.py
+1
-1
@@ -29,7 +29,7 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
29
from vllm.model_executor.models.interfaces import MixtureOfExperts
30
from vllm.model_executor.models.utils import maybe_prefix
31
from vllm.platforms import current_platform
32
-from vllm.utils import direct_register_custom_op
32
+from vllm.utils.torch_utils import direct_register_custom_op
33
34
from .utils import log_replacement
35
vllm/model_executor/models/utils.py
+4
-2
@@ -24,11 +24,13 @@ from vllm.multimodal import NestedTensors
24
from vllm.sequence import IntermediateTensors
25
from vllm.utils import (
26
cdiv,
27
- direct_register_custom_op,
28
- get_cuda_view_from_cpu_tensor,
27
is_pin_memory_available,
28
is_uva_available,
29
)
30
+from vllm.utils.torch_utils import (
31
+ direct_register_custom_op,
32
+ get_cuda_view_from_cpu_tensor,
33
+)
34
35
logger = init_logger(__name__)
36
vllm/model_executor/models/whisper.py
+1
-1
@@ -34,7 +34,6 @@ from vllm.model_executor.layers.linear import (
34
from vllm.model_executor.layers.logits_processor import LogitsProcessor
35
from vllm.model_executor.layers.quantization import QuantizationConfig
36
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
37
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.multimodal import MULTIMODAL_REGISTRY
39
from vllm.multimodal.inputs import (
@@ -53,6 +52,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
52
from vllm.transformers_utils.processor import cached_get_processor
53
from vllm.utils.jsontree import json_map_leaves
54
from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
+from vllm.utils.torch_utils import set_default_torch_dtype
56
57
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription
58
from .utils import (
vllm/model_executor/utils.py
+1
-1
@@ -7,7 +7,7 @@ from typing import Any
7
8
import torch
9
10
-from vllm.utils import is_torch_equal_or_newer
10
+from vllm.utils.torch_utils import is_torch_equal_or_newer
11
12
13
def set_random_seed(seed: int) -> None:
vllm/platforms/__init__.py
+1
-1
@@ -7,8 +7,8 @@ from typing import TYPE_CHECKING
7
8
from vllm import envs
9
from vllm.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group
10
-from vllm.utils import supports_xccl
10
from vllm.utils.import_utils import resolve_obj_by_qualname
11
+from vllm.utils.torch_utils import supports_xccl
12
13
from .interface import CpuArchEnum, Platform, PlatformEnum
14
vllm/platforms/cuda.py
+2
-1
@@ -16,7 +16,8 @@ from typing_extensions import ParamSpec
16
import vllm._C # noqa
17
import vllm.envs as envs
18
from vllm.logger import init_logger
19
-from vllm.utils import cuda_device_count_stateless, import_pynvml
19
+from vllm.utils import import_pynvml
20
+from vllm.utils.torch_utils import cuda_device_count_stateless
21
22
from .interface import DeviceCapability, Platform, PlatformEnum
23
vllm/platforms/rocm.py
+1
-1
@@ -9,7 +9,7 @@ import torch
9
10
import vllm.envs as envs
11
from vllm.logger import init_logger
12
-from vllm.utils import cuda_device_count_stateless
12
+from vllm.utils.torch_utils import cuda_device_count_stateless
13
14
from .interface import DeviceCapability, Platform, PlatformEnum
15
vllm/usage/usage_lib.py
+2
-1
@@ -21,7 +21,8 @@ import torch
21
import vllm.envs as envs
22
from vllm.connections import global_http_connection
23
from vllm.logger import init_logger
24
-from vllm.utils import cuda_device_count_stateless, cuda_get_device_properties
24
+from vllm.utils import cuda_get_device_properties
25
+from vllm.utils.torch_utils import cuda_device_count_stateless
26
from vllm.version import __version__ as VLLM_VERSION
27
28
logger = init_logger(__name__)
vllm/utils/__init__.py
+1
-577
@@ -35,12 +35,11 @@ from argparse import (
35
from collections import defaultdict
36
from collections.abc import (
37
Callable,
38
- Collection,
38
Iterator,
39
Sequence,
40
)
41
from concurrent.futures.process import ProcessPoolExecutor
43
-from functools import cache, lru_cache, partial, wraps
42
+from functools import cache, partial, wraps
43
from pathlib import Path
44
from typing import TYPE_CHECKING, Any, TextIO, TypeVar
45
from urllib.parse import urlparse
@@ -48,8 +47,6 @@ from uuid import uuid4
47
48
import cbor2
49
import cloudpickle
51
-import numpy as np
52
-import numpy.typing as npt
50
import psutil
51
import regex as re
52
import setproctitle
@@ -57,9 +54,6 @@ import torch
54
import yaml
55
import zmq
56
import zmq.asyncio
60
-from packaging import version
61
-from packaging.version import Version
62
-from torch.library import Library
57
58
import vllm.envs as envs
59
from vllm.logger import enable_trace_function_call, init_logger
@@ -69,13 +63,11 @@ if TYPE_CHECKING:
63
from argparse import Namespace
64
65
from vllm.config import ModelConfig, VllmConfig
72
- from vllm.sequence import IntermediateTensors
66
else:
67
Namespace = object
68
69
ModelConfig = object
70
VllmConfig = object
78
- IntermediateTensors = object
71
72
logger = init_logger(__name__)
73
@@ -105,46 +97,6 @@ STR_INVALID_VAL: str = "INVALID"
97
CYAN = "\033[1;36m"
98
RESET = "\033[0;0m"
99
108
-STR_DTYPE_TO_TORCH_DTYPE = {
109
- "float32": torch.float32,
110
- "half": torch.half,
111
- "bfloat16": torch.bfloat16,
112
- "float": torch.float,
113
- "fp8": torch.uint8,
114
- "fp8_e4m3": torch.uint8,
115
- "fp8_e5m2": torch.uint8,
116
- "int8": torch.int8,
117
- "fp8_inc": torch.float8_e4m3fn,
118
- "fp8_ds_mla": torch.uint8,
119
-}
120
-
121
-TORCH_DTYPE_TO_NUMPY_DTYPE = {
122
- torch.float16: np.float16,
123
- torch.float32: np.float32,
124
- torch.float64: np.float64,
125
- torch.uint8: np.uint8,
126
- torch.int32: np.int32,
127
- torch.int64: np.int64,
128
-}
129
-
130
-
131
-@contextlib.contextmanager
132
-def set_default_torch_num_threads(num_threads: int):
133
- """Sets the default number of threads for PyTorch to the given value."""
134
- old_num_threads = torch.get_num_threads()
135
- torch.set_num_threads(num_threads)
136
- yield
137
- torch.set_num_threads(old_num_threads)
138
-
139
-
140
-def kv_cache_dtype_str_to_dtype(
141
- kv_cache_dtype: str, model_config: ModelConfig
142
-) -> torch.dtype:
143
- if kv_cache_dtype == "auto":
144
- # Model config may not be specified for unit tests, default to float16
145
- return model_config.dtype if model_config else torch.half
146
- return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
147
-
100
101
T = TypeVar("T")
102
U = TypeVar("U")
@@ -407,141 +359,6 @@ def round_down(x: int, y: int) -> int:
359
return (x // y) * y
360
361
410
-def _generate_random_fp8(
411
- tensor: torch.Tensor,
412
- low: float,
413
- high: float,
414
-) -> None:
415
- # NOTE(zhaoyang): Due to NaN and Inf representation for fp8 data type,
416
- # it may occur Inf or NaN if we directly use torch.randint
417
- # to generate random data for fp8 data.
418
- # For example, s.11111.00 in fp8e5m2 format represents Inf.
419
- # | E4M3 | E5M2
420
- # -----|-------------|-------------------
421
- # Inf | N/A | s.11111.00
422
- # NaN | s.1111.111 | s.11111.{01,10,11}
423
- from vllm import _custom_ops as ops
424
-
425
- tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
426
- tensor_tmp.uniform_(low, high)
427
- ops.convert_fp8(tensor, tensor_tmp)
428
- del tensor_tmp
429
-
430
-
431
-def get_kv_cache_torch_dtype(
432
- cache_dtype: str | torch.dtype | None,
433
- model_dtype: str | torch.dtype | None = None,
434
-) -> torch.dtype:
435
- if isinstance(cache_dtype, str):
436
- if cache_dtype == "auto":
437
- if isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
438
- torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
439
- elif isinstance(model_dtype, torch.dtype):
440
- torch_dtype = model_dtype
441
- else:
442
- raise ValueError(f"Invalid model dtype: {model_dtype}")
443
- elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
444
- torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
445
- else:
446
- raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
447
- elif isinstance(cache_dtype, torch.dtype):
448
- torch_dtype = cache_dtype
449
- else:
450
- raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
451
- return torch_dtype
452
-
453
-
454
-def create_kv_caches_with_random_flash(
455
- num_blocks: int,
456
- block_size: int,
457
- num_layers: int,
458
- num_heads: int,
459
- head_size: int,
460
- cache_dtype: str | torch.dtype | None,
461
- model_dtype: str | torch.dtype | None = None,
462
- seed: int | None = None,
463
- device: str | None = "cuda",
464
- cache_layout: str | None = "NHD",
465
-) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
466
- from vllm.platforms import current_platform
467
-
468
- current_platform.seed_everything(seed)
469
-
470
- dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
471
- generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
472
- assert cache_layout in ("NHD", "HND")
473
- stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
474
-
475
- kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
476
- scale = head_size**-0.5
477
-
478
- key_caches: list[torch.Tensor] = []
479
- value_caches: list[torch.Tensor] = []
480
-
481
- for _ in range(num_layers):
482
- key_value_cache = torch.empty(
483
- size=kv_cache_allocation_shape, dtype=dtype, device=device
484
- ).permute(*stride_order)
485
- if cache_dtype in ["auto", "half", "bfloat16", "float"]:
486
- key_value_cache.uniform_(-scale, scale)
487
- elif cache_dtype == "fp8":
488
- _generate_random_fp8(key_value_cache, -scale, scale)
489
- else:
490
- raise ValueError(f"Does not support key cache of type {cache_dtype}")
491
- key_caches.append(key_value_cache[:, 0])
492
- value_caches.append(key_value_cache[:, 1])
493
- return key_caches, value_caches
494
-
495
-
496
-def create_kv_caches_with_random(
497
- num_blocks: int,
498
- block_size: int,
499
- num_layers: int,
500
- num_heads: int,
501
- head_size: int,
502
- cache_dtype: str | torch.dtype | None,
503
- model_dtype: str | torch.dtype | None = None,
504
- seed: int | None = None,
505
- device: str | None = "cuda",
506
-) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
507
- if cache_dtype == "fp8" and head_size % 16:
508
- raise ValueError(
509
- f"Does not support key cache of type fp8 with head_size {head_size}"
510
- )
511
- from vllm.platforms import current_platform
512
-
513
- current_platform.seed_everything(seed)
514
-
515
- dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
516
-
517
- scale = head_size**-0.5
518
- x = 16 // torch.tensor([], dtype=dtype).element_size()
519
- key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
520
- key_caches: list[torch.Tensor] = []
521
- for _ in range(num_layers):
522
- key_cache = torch.empty(size=key_cache_shape, dtype=dtype, device=device)
523
- if cache_dtype in ["auto", "half", "bfloat16", "float"]:
524
- key_cache.uniform_(-scale, scale)
525
- elif cache_dtype == "fp8":
526
- _generate_random_fp8(key_cache, -scale, scale)
527
- else:
528
- raise ValueError(f"Does not support key cache of type {cache_dtype}")
529
- key_caches.append(key_cache)
530
-
531
- value_cache_shape = (num_blocks, num_heads, head_size, block_size)
532
- value_caches: list[torch.Tensor] = []
533
- for _ in range(num_layers):
534
- value_cache = torch.empty(size=value_cache_shape, dtype=dtype, device=device)
535
- if cache_dtype in ["auto", "half", "bfloat16", "float"]:
536
- value_cache.uniform_(-scale, scale)
537
- elif cache_dtype == "fp8":
538
- _generate_random_fp8(value_cache, -scale, scale)
539
- else:
540
- raise ValueError(f"Does not support value cache of type {cache_dtype}")
541
- value_caches.append(value_cache)
542
- return key_caches, value_caches
543
-
544
-
362
@cache
363
def is_pin_memory_available() -> bool:
364
from vllm.platforms import current_platform
@@ -557,121 +374,6 @@ def is_uva_available() -> bool:
374
return is_pin_memory_available()
375
376
560
-def make_ndarray_with_pad(
561
- x: list[list[T]],
562
- pad: T,
563
- dtype: npt.DTypeLike,
564
- *,
565
- max_len: int | None = None,
566
-) -> npt.NDArray:
567
- """
568
- Make a padded array from 2D inputs.
569
-
570
- The padding is applied to the end of each inner list until it reaches
571
- `max_len`.
572
- """
573
- if max_len is None:
574
- # Unlike for most functions, map is faster than a genexpr over `len`
575
- max_len = max(map(len, x), default=0)
576
-
577
- padded_x = np.full((len(x), max_len), pad, dtype=dtype)
578
- for ind, blocktb in enumerate(x):
579
- assert len(blocktb) <= max_len
580
- padded_x[ind, : len(blocktb)] = blocktb
581
-
582
- return padded_x
583
-
584
-
585
-def make_tensor_with_pad(
586
- x: list[list[T]],
587
- pad: T,
588
- dtype: torch.dtype,
589
- *,
590
- max_len: int | None = None,
591
- device: str | torch.device | None = None,
592
- pin_memory: bool = False,
593
-) -> torch.Tensor:
594
- """
595
- Make a padded tensor from 2D inputs.
596
-
597
- The padding is applied to the end of each inner list until it reaches
598
- `max_len`.
599
- """
600
- np_dtype = TORCH_DTYPE_TO_NUMPY_DTYPE[dtype]
601
- padded_x = make_ndarray_with_pad(x, pad, np_dtype, max_len=max_len)
602
-
603
- tensor = torch.from_numpy(padded_x).to(device)
604
- if pin_memory:
605
- tensor = tensor.pin_memory()
606
-
607
- return tensor
608
-
609
-
610
-def async_tensor_h2d(
611
- data: list,
612
- dtype: torch.dtype,
613
- target_device: str | torch.device,
614
- pin_memory: bool,
615
-) -> torch.Tensor:
616
- """Asynchronously create a tensor and copy it from host to device."""
617
- t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")
618
- return t.to(device=target_device, non_blocking=True)
619
-
620
-
621
-def get_dtype_size(dtype: torch.dtype) -> int:
622
- """Get the size of the data type in bytes."""
623
- return torch.tensor([], dtype=dtype).element_size()
624
-
625
-
626
-# bool = 0, int = 1, float = 2, complex = 3
627
-def _get_precision_level(dtype: torch.dtype) -> int:
628
- # NOTE: Complex dtypes return `is_floating_point=False`
629
- return (dtype != torch.bool) + dtype.is_floating_point + dtype.is_complex * 2
630
-
631
-
632
-def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):
633
- """
634
- Test whether it is lossless to cast a tensor from
635
- `src_dtype` to `tgt_dtype`.
636
- """
637
- if src_dtype == tgt_dtype:
638
- return True
639
-
640
- src_level = _get_precision_level(src_dtype)
641
- tgt_level = _get_precision_level(tgt_dtype)
642
-
643
- if src_level < tgt_level:
644
- return True
645
- if src_level > tgt_level:
646
- return False
647
-
648
- # Compare integral types
649
- if not src_dtype.is_floating_point and not src_dtype.is_complex:
650
- src_info = torch.iinfo(src_dtype)
651
- tgt_info = torch.iinfo(tgt_dtype)
652
- return src_info.min >= tgt_info.min and src_info.max <= tgt_info.max
653
-
654
- # Compare floating-point types
655
- src_info = torch.finfo(src_dtype)
656
- tgt_info = torch.finfo(tgt_dtype)
657
- return (
658
- src_info.min >= tgt_info.min
659
- and src_info.max <= tgt_info.max
660
- and src_info.resolution >= tgt_info.resolution
661
- )
662
-
663
-
664
-def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
665
- """
666
- Get the common `dtype` where all of the other `dtypes` can be
667
- cast to it without losing any information.
668
- """
669
- return max(
670
- dtypes,
671
- key=lambda dtype: sum(is_lossless_cast(dt, dtype) for dt in dtypes),
672
- )
673
-
674
-
377
# TODO: This function can be removed if transformer_modules classes are
378
# serialized by value when communicating between processes
379
def init_cached_hf_modules() -> None:
@@ -767,60 +469,6 @@ def find_nccl_include_paths() -> list[str] | None:
469
return out or None
470
471
770
-prev_set_stream = torch.cuda.set_stream
771
-
772
-_current_stream_tls = threading.local()
773
-
774
-
775
-def _patched_set_stream(stream: torch.cuda.Stream) -> None:
776
- _current_stream_tls.value = stream
777
- prev_set_stream(stream)
778
-
779
-
780
-torch.cuda.set_stream = _patched_set_stream
781
-
782
-
783
-class _StreamPlaceholder:
784
- def __init__(self):
785
- self.synchronize = lambda: None
786
-
787
-
788
-def current_stream() -> torch.cuda.Stream:
789
- """
790
- replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
791
- it turns out that `torch.cuda.current_stream()` is quite expensive,
792
- as it will construct a new stream object at each call.
793
- here we patch `torch.cuda.set_stream` to keep track of the current stream
794
- directly, so that we can avoid calling `torch.cuda.current_stream()`.
795
-
796
- the underlying hypothesis is that we do not call `torch._C._cuda_setStream`
797
- from C/C++ code.
798
- """
799
- from vllm.platforms import current_platform
800
-
801
- if not hasattr(_current_stream_tls, "value") or _current_stream_tls.value is None:
802
- # when this function is called before any stream is set,
803
- # we return the default stream.
804
- # On ROCm using the default 0 stream in combination with RCCL
805
- # is hurting performance. Therefore creating a dedicated stream
806
- # per process
807
- if current_platform.is_rocm():
808
- # torch.cuda.set_stream here is the alias of _pathed_set_stream
809
- torch.cuda.set_stream(torch.cuda.Stream())
810
- elif current_platform.is_cpu():
811
- _current_stream_tls.value = _StreamPlaceholder()
812
- else:
813
- current_stream = current_platform.current_stream
814
- if current_stream is not None:
815
- _current_stream_tls.value = current_stream()
816
- else:
817
- raise ValueError(
818
- "Fail to set current stream, current platform "
819
- "may not support current_stream with torch API"
820
- )
821
- return _current_stream_tls.value
822
-
823
-
472
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
473
"""Set up function tracing for the current thread,
474
if enabled via the VLLM_TRACE_FUNCTION environment variable
@@ -842,48 +490,6 @@ def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
490
enable_trace_function_call(log_path)
491
492
845
-@lru_cache(maxsize=8)
846
-def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:
847
- # Note: cuda_visible_devices is not used, but we keep it as an argument for
848
- # LRU Cache purposes.
849
-
850
- # Code below is based on
851
- # https://github.com/pytorch/pytorch/blob/
852
- # c1cd946818442aca8c7f812b16d187ce1586c3bc/
853
- # torch/cuda/__init__.py#L831C1-L831C17
854
- import torch.cuda
855
-
856
- from vllm.platforms import current_platform
857
-
858
- if not torch.cuda._is_compiled():
859
- return 0
860
- if current_platform.is_rocm():
861
- # ROCm uses amdsmi instead of nvml for stateless device count
862
- # This requires a sufficiently modern version of Torch 2.4.0
863
- raw_count = (
864
- torch.cuda._device_count_amdsmi()
865
- if (hasattr(torch.cuda, "_device_count_amdsmi"))
866
- else -1
867
- )
868
- else:
869
- raw_count = torch.cuda._device_count_nvml()
870
- r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
871
- return r
872
-
873
-
874
-def cuda_device_count_stateless() -> int:
875
- """Get number of CUDA devices, caching based on the value of
876
- CUDA_VISIBLE_DEVICES at the time of call.
877
-
878
- This should be used instead of torch.cuda.device_count()
879
- unless CUDA_VISIBLE_DEVICES has already been set to the desired
880
- value."""
881
-
882
- # This can be removed and simply replaced with torch.cuda.get_device_count
883
- # after https://github.com/pytorch/pytorch/pull/122815 is released.
884
- return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)
885
-
886
-
493
def cuda_is_initialized() -> bool:
494
"""Check if CUDA is initialized."""
495
if not torch.cuda._is_compiled():
@@ -1411,27 +1017,6 @@ class FlexibleArgumentParser(ArgumentParser):
1017
return processed_args
1018
1019
1414
-# Using dynamo with vLLM doesn't really work well with PyTorch versions < 2.4.0.
1415
-# In particular, the FakeScalarType is not supported for earlier versions of
1416
-# PyTorch which breaks dynamo for any ops registered using ScalarType.
1417
-def supports_dynamo() -> bool:
1418
- base_torch_version = Version(Version(torch.__version__).base_version)
1419
- return base_torch_version >= Version("2.4.0")
1420
-
1421
-
1422
-# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
1423
-def supports_xccl() -> bool:
1424
- return (
1425
- is_torch_equal_or_newer("2.8.0.dev") and torch.distributed.is_xccl_available()
1426
- )
1427
-
1428
-
1429
-# Some backends use pytorch version < 2.4.0 which doesn't
1430
-# support `torch.library.custom_op`.
1431
-def supports_custom_op() -> bool:
1432
- return hasattr(torch.library, "custom_op")
1433
-
1434
-
1020
class AtomicCounter:
1021
"""An atomic, thread-safe counter"""
1022
@@ -1457,118 +1042,6 @@ class AtomicCounter:
1042
return self._value
1043
1044
1460
-def weak_ref_tensor(tensor: Any) -> Any:
1461
- """
1462
- Create a weak reference to a tensor.
1463
- The new tensor will share the same data as the original tensor,
1464
- but will not keep the original tensor alive.
1465
- """
1466
- if isinstance(tensor, torch.Tensor):
1467
- return torch.ops._C.weak_ref_tensor(tensor)
1468
- else:
1469
- return tensor
1470
-
1471
-
1472
-def weak_ref_tensors(
1473
- tensors: torch.Tensor
1474
- | list[torch.Tensor]
1475
- | tuple[torch.Tensor]
1476
- | IntermediateTensors,
1477
-) -> torch.Tensor | list[Any] | tuple[Any] | Any:
1478
- """
1479
- Convenience function to create weak references to tensors,
1480
- for single tensor, list of tensors or tuple of tensors.
1481
- """
1482
- if isinstance(tensors, torch.Tensor):
1483
- return weak_ref_tensor(tensors)
1484
- if isinstance(tensors, list):
1485
- return [weak_ref_tensor(t) for t in tensors]
1486
- if isinstance(tensors, tuple):
1487
- return tuple(weak_ref_tensor(t) for t in tensors)
1488
-
1489
- # For IntermediateTensors used in pipeline parallelism
1490
- from vllm.sequence import IntermediateTensors
1491
-
1492
- if isinstance(tensors, IntermediateTensors):
1493
- ret = IntermediateTensors(
1494
- {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}
1495
- )
1496
- return ret
1497
- raise ValueError("Invalid type for tensors")
1498
-
1499
-
1500
-def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
1501
- """
1502
- Get a CUDA view of a CPU tensor using Unified Virtual Addressing (UVA).
1503
- """
1504
- assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
1505
- return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
1506
-
1507
-
1508
-# create a library to hold the custom op
1509
-vllm_lib = Library("vllm", "FRAGMENT") # noqa
1510
-
1511
-
1512
-def direct_register_custom_op(
1513
- op_name: str,
1514
- op_func: Callable,
1515
- mutates_args: list[str] | None = None,
1516
- fake_impl: Callable | None = None,
1517
- target_lib: Library | None = None,
1518
- dispatch_key: str | None = None,
1519
- tags: tuple[torch.Tag, ...] = (),
1520
-):
1521
- """
1522
- `torch.library.custom_op` can have significant overhead because it
1523
- needs to consider complicated dispatching logic. This function
1524
- directly registers a custom op and dispatches it to the CUDA backend.
1525
- See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
1526
- for more details.
1527
-
1528
- By default, the custom op is registered to the vLLM library. If you
1529
- want to register it to a different library, you can pass the library
1530
- object to the `target_lib` argument.
1531
-
1532
- IMPORTANT: the lifetime of the operator is tied to the lifetime of the
1533
- library object. If you want to bind the operator to a different library,
1534
- make sure the library object is alive when the operator is used.
1535
- """
1536
- if not supports_custom_op():
1537
- from vllm.platforms import current_platform
1538
-
1539
- assert not current_platform.is_cuda_alike(), (
1540
- "cuda platform needs torch>=2.4 to support custom op, "
1541
- "chances are you are using an old version of pytorch "
1542
- "or a custom build of pytorch. It is recommended to "
1543
- "use vLLM in a fresh new environment and let it install "
1544
- "the required dependencies."
1545
- )
1546
- return
1547
-
1548
- if mutates_args is None:
1549
- mutates_args = []
1550
-
1551
- if dispatch_key is None:
1552
- from vllm.platforms import current_platform
1553
-
1554
- dispatch_key = current_platform.dispatch_key
1555
-
1556
- import torch.library
1557
-
1558
- if hasattr(torch.library, "infer_schema"):
1559
- schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
1560
- else:
1561
- # for pytorch 2.4
1562
- import torch._custom_op.impl
1563
-
1564
- schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
1565
- my_lib = target_lib or vllm_lib
1566
- my_lib.define(op_name + schema_str, tags=tags)
1567
- my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
1568
- if fake_impl is not None:
1569
- my_lib._register_fake(op_name, fake_impl)
1570
-
1571
-
1045
def kill_process_tree(pid: int):
1046
"""
1047
Kills all descendant processes of the given pid by sending SIGKILL.
@@ -2063,55 +1536,6 @@ def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
1536
raise ValueError(f"Unsupported hash function: {hash_fn_name}")
1537
1538
2066
-def is_torch_equal_or_newer(target: str) -> bool:
2067
- """Check if the installed torch version is >= the target version.
2068
-
2069
- Args:
2070
- target: a version string, like "2.6.0".
2071
-
2072
- Returns:
2073
- Whether the condition meets.
2074
- """
2075
- try:
2076
- return _is_torch_equal_or_newer(str(torch.__version__), target)
2077
- except Exception:
2078
- # Fallback to PKG-INFO to load the package info, needed by the doc gen.
2079
- return Version(importlib.metadata.version("torch")) >= Version(target)
2080
-
2081
-
2082
-# Helper function used in testing.
2083
-def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool:
2084
- torch_version = version.parse(torch_version)
2085
- return torch_version >= version.parse(target)
2086
-
2087
-
2088
-def _is_torch_equal(target: str) -> bool:
2089
- assert target.count(".") == 2
2090
- torch_version = str(torch.__version__)
2091
- torch_version = version.parse(torch_version)
2092
- # torch version is like "2.6.0.dev20240101" or "2.6.0.dev20240101+cpu"
2093
- # or "2.6.0+cu128" but never "2.6.0.1"
2094
- return (
2095
- torch_version >= version.parse(target)
2096
- and version.parse(target + ".1") > torch_version
2097
- )
2098
-
2099
-
2100
-def is_torch_equal(target: str) -> bool:
2101
- """Check if the installed torch version is == the target version.
2102
-
2103
- Args:
2104
- target: a version string, like "2.6.0".
2105
-
2106
- Returns:
2107
- Whether the condition meets.
2108
- """
2109
- try:
2110
- return _is_torch_equal(target)
2111
- except Exception:
2112
- return Version(importlib.metadata.version("torch")) == Version(target)
2113
-
2114
-
1539
@cache
1540
def _has_module(module_name: str) -> bool:
1541
"""Return True if *module_name* can be found in the current environment.
vllm/utils/torch_utils.py
new
+605
@@ -0,0 +1,605 @@
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 threading
6
+from collections.abc import Callable, Collection
7
+from functools import lru_cache
8
+from typing import TYPE_CHECKING, Any, TypeVar
9
+
10
+import numpy as np
11
+import numpy.typing as npt
12
+import torch
13
+from packaging import version
14
+from packaging.version import Version
15
+from torch.library import Library
16
+
17
+import vllm.envs as envs
18
+
19
+if TYPE_CHECKING:
20
+ from vllm.config import ModelConfig
21
+ from vllm.sequence import IntermediateTensors
22
+else:
23
+ ModelConfig = object
24
+ IntermediateTensors = object
25
+
26
+
27
+STR_DTYPE_TO_TORCH_DTYPE = {
28
+ "float32": torch.float32,
29
+ "half": torch.half,
30
+ "bfloat16": torch.bfloat16,
31
+ "float": torch.float,
32
+ "fp8": torch.uint8,
33
+ "fp8_e4m3": torch.uint8,
34
+ "fp8_e5m2": torch.uint8,
35
+ "int8": torch.int8,
36
+ "fp8_inc": torch.float8_e4m3fn,
37
+ "fp8_ds_mla": torch.uint8,
38
+}
39
+
40
+TORCH_DTYPE_TO_NUMPY_DTYPE = {
41
+ torch.float16: np.float16,
42
+ torch.float32: np.float32,
43
+ torch.float64: np.float64,
44
+ torch.uint8: np.uint8,
45
+ torch.int32: np.int32,
46
+ torch.int64: np.int64,
47
+}
48
+
49
+
50
+T = TypeVar("T")
51
+
52
+
53
+@contextlib.contextmanager
54
+def set_default_torch_dtype(dtype: torch.dtype):
55
+ """Sets the default torch dtype to the given dtype."""
56
+ old_dtype = torch.get_default_dtype()
57
+ torch.set_default_dtype(dtype)
58
+ yield
59
+ torch.set_default_dtype(old_dtype)
60
+
61
+
62
+@contextlib.contextmanager
63
+def set_default_torch_num_threads(num_threads: int):
64
+ """Sets the default number of threads for PyTorch to the given value."""
65
+ old_num_threads = torch.get_num_threads()
66
+ torch.set_num_threads(num_threads)
67
+ yield
68
+ torch.set_num_threads(old_num_threads)
69
+
70
+
71
+def get_dtype_size(dtype: torch.dtype) -> int:
72
+ """Get the size of the data type in bytes."""
73
+ return torch.tensor([], dtype=dtype).element_size()
74
+
75
+
76
+# bool = 0, int = 1, float = 2, complex = 3
77
+def _get_precision_level(dtype: torch.dtype) -> int:
78
+ # NOTE: Complex dtypes return `is_floating_point=False`
79
+ return (dtype != torch.bool) + dtype.is_floating_point + dtype.is_complex * 2
80
+
81
+
82
+def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):
83
+ """
84
+ Test whether it is lossless to cast a tensor from
85
+ `src_dtype` to `tgt_dtype`.
86
+ """
87
+ if src_dtype == tgt_dtype:
88
+ return True
89
+
90
+ src_level = _get_precision_level(src_dtype)
91
+ tgt_level = _get_precision_level(tgt_dtype)
92
+
93
+ if src_level < tgt_level:
94
+ return True
95
+ if src_level > tgt_level:
96
+ return False
97
+
98
+ # Compare integral types
99
+ if not src_dtype.is_floating_point and not src_dtype.is_complex:
100
+ src_info = torch.iinfo(src_dtype)
101
+ tgt_info = torch.iinfo(tgt_dtype)
102
+ return src_info.min >= tgt_info.min and src_info.max <= tgt_info.max
103
+
104
+ # Compare floating-point types
105
+ src_info = torch.finfo(src_dtype)
106
+ tgt_info = torch.finfo(tgt_dtype)
107
+ return (
108
+ src_info.min >= tgt_info.min
109
+ and src_info.max <= tgt_info.max
110
+ and src_info.resolution >= tgt_info.resolution
111
+ )
112
+
113
+
114
+def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
115
+ """
116
+ Get the common `dtype` where all of the other `dtypes` can be
117
+ cast to it without losing any information.
118
+ """
119
+ return max(
120
+ dtypes,
121
+ key=lambda dtype: sum(is_lossless_cast(dt, dtype) for dt in dtypes),
122
+ )
123
+
124
+
125
+def _generate_random_fp8(
126
+ tensor: torch.Tensor,
127
+ low: float,
128
+ high: float,
129
+) -> None:
130
+ # NOTE(zhaoyang): Due to NaN and Inf representation for fp8 data type,
131
+ # it may occur Inf or NaN if we directly use torch.randint
132
+ # to generate random data for fp8 data.
133
+ # For example, s.11111.00 in fp8e5m2 format represents Inf.
134
+ # | E4M3 | E5M2
135
+ # -----|-------------|-------------------
136
+ # Inf | N/A | s.11111.00
137
+ # NaN | s.1111.111 | s.11111.{01,10,11}
138
+ from vllm import _custom_ops as ops
139
+
140
+ tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
141
+ tensor_tmp.uniform_(low, high)
142
+ ops.convert_fp8(tensor, tensor_tmp)
143
+ del tensor_tmp
144
+
145
+
146
+def get_kv_cache_torch_dtype(
147
+ cache_dtype: str | torch.dtype | None,
148
+ model_dtype: str | torch.dtype | None = None,
149
+) -> torch.dtype:
150
+ if isinstance(cache_dtype, str):
151
+ if cache_dtype == "auto":
152
+ if isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
153
+ torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
154
+ elif isinstance(model_dtype, torch.dtype):
155
+ torch_dtype = model_dtype
156
+ else:
157
+ raise ValueError(f"Invalid model dtype: {model_dtype}")
158
+ elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
159
+ torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
160
+ else:
161
+ raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
162
+ elif isinstance(cache_dtype, torch.dtype):
163
+ torch_dtype = cache_dtype
164
+ else:
165
+ raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
166
+ return torch_dtype
167
+
168
+
169
+def kv_cache_dtype_str_to_dtype(
170
+ kv_cache_dtype: str, model_config: ModelConfig
171
+) -> torch.dtype:
172
+ if kv_cache_dtype == "auto":
173
+ # Model config may not be specified for unit tests, default to float16
174
+ return model_config.dtype if model_config else torch.half
175
+ return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
176
+
177
+
178
+def create_kv_caches_with_random_flash(
179
+ num_blocks: int,
180
+ block_size: int,
181
+ num_layers: int,
182
+ num_heads: int,
183
+ head_size: int,
184
+ cache_dtype: str | torch.dtype | None,
185
+ model_dtype: str | torch.dtype | None = None,
186
+ seed: int | None = None,
187
+ device: str | None = "cuda",
188
+ cache_layout: str | None = "NHD",
189
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
190
+ from vllm.platforms import current_platform
191
+
192
+ current_platform.seed_everything(seed)
193
+
194
+ dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
195
+ generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
196
+ assert cache_layout in ("NHD", "HND")
197
+ stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
198
+
199
+ kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
200
+ scale = head_size**-0.5
201
+
202
+ key_caches: list[torch.Tensor] = []
203
+ value_caches: list[torch.Tensor] = []
204
+
205
+ for _ in range(num_layers):
206
+ key_value_cache = torch.empty(
207
+ size=kv_cache_allocation_shape, dtype=dtype, device=device
208
+ ).permute(*stride_order)
209
+ if cache_dtype in ["auto", "half", "bfloat16", "float"]:
210
+ key_value_cache.uniform_(-scale, scale)
211
+ elif cache_dtype == "fp8":
212
+ _generate_random_fp8(key_value_cache, -scale, scale)
213
+ else:
214
+ raise ValueError(f"Does not support key cache of type {cache_dtype}")
215
+ key_caches.append(key_value_cache[:, 0])
216
+ value_caches.append(key_value_cache[:, 1])
217
+ return key_caches, value_caches
218
+
219
+
220
+def create_kv_caches_with_random(
221
+ num_blocks: int,
222
+ block_size: int,
223
+ num_layers: int,
224
+ num_heads: int,
225
+ head_size: int,
226
+ cache_dtype: str | torch.dtype | None,
227
+ model_dtype: str | torch.dtype | None = None,
228
+ seed: int | None = None,
229
+ device: str | None = "cuda",
230
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
231
+ if cache_dtype == "fp8" and head_size % 16:
232
+ raise ValueError(
233
+ f"Does not support key cache of type fp8 with head_size {head_size}"
234
+ )
235
+ from vllm.platforms import current_platform
236
+
237
+ current_platform.seed_everything(seed)
238
+
239
+ dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
240
+
241
+ scale = head_size**-0.5
242
+ x = 16 // torch.tensor([], dtype=dtype).element_size()
243
+ key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
244
+ key_caches: list[torch.Tensor] = []
245
+ for _ in range(num_layers):
246
+ key_cache = torch.empty(size=key_cache_shape, dtype=dtype, device=device)
247
+ if cache_dtype in ["auto", "half", "bfloat16", "float"]:
248
+ key_cache.uniform_(-scale, scale)
249
+ elif cache_dtype == "fp8":
250
+ _generate_random_fp8(key_cache, -scale, scale)
251
+ else:
252
+ raise ValueError(f"Does not support key cache of type {cache_dtype}")
253
+ key_caches.append(key_cache)
254
+
255
+ value_cache_shape = (num_blocks, num_heads, head_size, block_size)
256
+ value_caches: list[torch.Tensor] = []
257
+ for _ in range(num_layers):
258
+ value_cache = torch.empty(size=value_cache_shape, dtype=dtype, device=device)
259
+ if cache_dtype in ["auto", "half", "bfloat16", "float"]:
260
+ value_cache.uniform_(-scale, scale)
261
+ elif cache_dtype == "fp8":
262
+ _generate_random_fp8(value_cache, -scale, scale)
263
+ else:
264
+ raise ValueError(f"Does not support value cache of type {cache_dtype}")
265
+ value_caches.append(value_cache)
266
+ return key_caches, value_caches
267
+
268
+
269
+def async_tensor_h2d(
270
+ data: list,
271
+ dtype: torch.dtype,
272
+ target_device: str | torch.device,
273
+ pin_memory: bool,
274
+) -> torch.Tensor:
275
+ """Asynchronously create a tensor and copy it from host to device."""
276
+ t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")
277
+ return t.to(device=target_device, non_blocking=True)
278
+
279
+
280
+def make_ndarray_with_pad(
281
+ x: list[list[T]],
282
+ pad: T,
283
+ dtype: npt.DTypeLike,
284
+ *,
285
+ max_len: int | None = None,
286
+) -> npt.NDArray:
287
+ """
288
+ Make a padded array from 2D inputs.
289
+
290
+ The padding is applied to the end of each inner list until it reaches
291
+ `max_len`.
292
+ """
293
+ if max_len is None:
294
+ # Unlike for most functions, map is faster than a genexpr over `len`
295
+ max_len = max(map(len, x), default=0)
296
+
297
+ padded_x = np.full((len(x), max_len), pad, dtype=dtype)
298
+ for ind, blocktb in enumerate(x):
299
+ assert len(blocktb) <= max_len
300
+ padded_x[ind, : len(blocktb)] = blocktb
301
+
302
+ return padded_x
303
+
304
+
305
+def make_tensor_with_pad(
306
+ x: list[list[T]],
307
+ pad: T,
308
+ dtype: torch.dtype,
309
+ *,
310
+ max_len: int | None = None,
311
+ device: str | torch.device | None = None,
312
+ pin_memory: bool = False,
313
+) -> torch.Tensor:
314
+ """
315
+ Make a padded tensor from 2D inputs.
316
+
317
+ The padding is applied to the end of each inner list until it reaches
318
+ `max_len`.
319
+ """
320
+ np_dtype = TORCH_DTYPE_TO_NUMPY_DTYPE[dtype]
321
+ padded_x = make_ndarray_with_pad(x, pad, np_dtype, max_len=max_len)
322
+
323
+ tensor = torch.from_numpy(padded_x).to(device)
324
+ if pin_memory:
325
+ tensor = tensor.pin_memory()
326
+
327
+ return tensor
328
+
329
+
330
+prev_set_stream = torch.cuda.set_stream
331
+
332
+_current_stream_tls = threading.local()
333
+
334
+
335
+def _patched_set_stream(stream: torch.cuda.Stream) -> None:
336
+ _current_stream_tls.value = stream
337
+ prev_set_stream(stream)
338
+
339
+
340
+torch.cuda.set_stream = _patched_set_stream
341
+
342
+
343
+class _StreamPlaceholder:
344
+ def __init__(self):
345
+ self.synchronize = lambda: None
346
+
347
+
348
+def current_stream() -> torch.cuda.Stream:
349
+ """
350
+ replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
351
+ it turns out that `torch.cuda.current_stream()` is quite expensive,
352
+ as it will construct a new stream object at each call.
353
+ here we patch `torch.cuda.set_stream` to keep track of the current stream
354
+ directly, so that we can avoid calling `torch.cuda.current_stream()`.
355
+
356
+ the underlying hypothesis is that we do not call `torch._C._cuda_setStream`
357
+ from C/C++ code.
358
+ """
359
+ from vllm.platforms import current_platform
360
+
361
+ if not hasattr(_current_stream_tls, "value") or _current_stream_tls.value is None:
362
+ # when this function is called before any stream is set,
363
+ # we return the default stream.
364
+ # On ROCm using the default 0 stream in combination with RCCL
365
+ # is hurting performance. Therefore creating a dedicated stream
366
+ # per process
367
+ if current_platform.is_rocm():
368
+ # torch.cuda.set_stream here is the alias of _pathed_set_stream
369
+ torch.cuda.set_stream(torch.cuda.Stream())
370
+ elif current_platform.is_cpu():
371
+ _current_stream_tls.value = _StreamPlaceholder()
372
+ else:
373
+ current_stream = current_platform.current_stream
374
+ if current_stream is not None:
375
+ _current_stream_tls.value = current_stream()
376
+ else:
377
+ raise ValueError(
378
+ "Fail to set current stream, current platform "
379
+ "may not support current_stream with torch API"
380
+ )
381
+ return _current_stream_tls.value
382
+
383
+
384
+@lru_cache(maxsize=8)
385
+def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:
386
+ # Note: cuda_visible_devices is not used, but we keep it as an argument for
387
+ # LRU Cache purposes.
388
+
389
+ # Code below is based on
390
+ # https://github.com/pytorch/pytorch/blob/
391
+ # c1cd946818442aca8c7f812b16d187ce1586c3bc/
392
+ # torch/cuda/__init__.py#L831C1-L831C17
393
+ import torch.cuda
394
+ import torch.version
395
+
396
+ from vllm.platforms import current_platform
397
+
398
+ if not torch.cuda._is_compiled():
399
+ return 0
400
+ if current_platform.is_rocm():
401
+ # ROCm uses amdsmi instead of nvml for stateless device count
402
+ # This requires a sufficiently modern version of Torch 2.4.0
403
+ raw_count = (
404
+ torch.cuda._device_count_amdsmi()
405
+ if (hasattr(torch.cuda, "_device_count_amdsmi"))
406
+ else -1
407
+ )
408
+ else:
409
+ raw_count = torch.cuda._device_count_nvml()
410
+ r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
411
+ return r
412
+
413
+
414
+def cuda_device_count_stateless() -> int:
415
+ """Get number of CUDA devices, caching based on the value of
416
+ CUDA_VISIBLE_DEVICES at the time of call.
417
+
418
+ This should be used instead of torch.cuda.device_count()
419
+ unless CUDA_VISIBLE_DEVICES has already been set to the desired
420
+ value."""
421
+
422
+ # This can be removed and simply replaced with torch.cuda.get_device_count
423
+ # after https://github.com/pytorch/pytorch/pull/122815 is released.
424
+ return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)
425
+
426
+
427
+def weak_ref_tensor(tensor: Any) -> Any:
428
+ """
429
+ Create a weak reference to a tensor.
430
+ The new tensor will share the same data as the original tensor,
431
+ but will not keep the original tensor alive.
432
+ """
433
+ if isinstance(tensor, torch.Tensor):
434
+ return torch.ops._C.weak_ref_tensor(tensor)
435
+ else:
436
+ return tensor
437
+
438
+
439
+def weak_ref_tensors(
440
+ tensors: torch.Tensor
441
+ | list[torch.Tensor]
442
+ | tuple[torch.Tensor]
443
+ | IntermediateTensors,
444
+) -> torch.Tensor | list[Any] | tuple[Any] | Any:
445
+ """
446
+ Convenience function to create weak references to tensors,
447
+ for single tensor, list of tensors or tuple of tensors.
448
+ """
449
+ if isinstance(tensors, torch.Tensor):
450
+ return weak_ref_tensor(tensors)
451
+ if isinstance(tensors, list):
452
+ return [weak_ref_tensor(t) for t in tensors]
453
+ if isinstance(tensors, tuple):
454
+ return tuple(weak_ref_tensor(t) for t in tensors)
455
+
456
+ # For IntermediateTensors used in pipeline parallelism
457
+ from vllm.sequence import IntermediateTensors
458
+
459
+ if isinstance(tensors, IntermediateTensors):
460
+ ret = IntermediateTensors(
461
+ {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}
462
+ )
463
+ return ret
464
+ raise ValueError("Invalid type for tensors")
465
+
466
+
467
+def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
468
+ """
469
+ Get a CUDA view of a CPU tensor using Unified Virtual Addressing (UVA).
470
+ """
471
+ assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
472
+ return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
473
+
474
+
475
+# Helper function used in testing.
476
+def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool:
477
+ torch_version = version.parse(torch_version)
478
+ return torch_version >= version.parse(target)
479
+
480
+
481
+def is_torch_equal_or_newer(target: str) -> bool:
482
+ """Check if the installed torch version is >= the target version.
483
+
484
+ Args:
485
+ target: a version string, like "2.6.0".
486
+
487
+ Returns:
488
+ Whether the condition meets.
489
+ """
490
+ try:
491
+ return _is_torch_equal_or_newer(str(torch.__version__), target)
492
+ except Exception:
493
+ # Fallback to PKG-INFO to load the package info, needed by the doc gen.
494
+ return Version(importlib.metadata.version("torch")) >= Version(target)
495
+
496
+
497
+def _is_torch_equal(target: str) -> bool:
498
+ assert target.count(".") == 2
499
+ torch_version = str(torch.__version__)
500
+ torch_version = version.parse(torch_version)
501
+ # torch version is like "2.6.0.dev20240101" or "2.6.0.dev20240101+cpu"
502
+ # or "2.6.0+cu128" but never "2.6.0.1"
503
+ return (
504
+ torch_version >= version.parse(target)
505
+ and version.parse(target + ".1") > torch_version
506
+ )
507
+
508
+
509
+def is_torch_equal(target: str) -> bool:
510
+ """Check if the installed torch version is == the target version.
511
+
512
+ Args:
513
+ target: a version string, like "2.6.0".
514
+
515
+ Returns:
516
+ Whether the condition meets.
517
+ """
518
+ try:
519
+ return _is_torch_equal(target)
520
+ except Exception:
521
+ return Version(importlib.metadata.version("torch")) == Version(target)
522
+
523
+
524
+# Using dynamo with vLLM doesn't really work well with PyTorch versions < 2.4.0.
525
+# In particular, the FakeScalarType is not supported for earlier versions of
526
+# PyTorch which breaks dynamo for any ops registered using ScalarType.
527
+def supports_dynamo() -> bool:
528
+ return is_torch_equal_or_newer("2.4.0")
529
+
530
+
531
+# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
532
+def supports_xccl() -> bool:
533
+ return (
534
+ is_torch_equal_or_newer("2.8.0.dev") and torch.distributed.is_xccl_available()
535
+ )
536
+
537
+
538
+# Some backends use pytorch version < 2.4.0 which doesn't
539
+# support `torch.library.custom_op`.
540
+def supports_custom_op() -> bool:
541
+ return hasattr(torch.library, "custom_op")
542
+
543
+
544
+# create a library to hold the custom op
545
+vllm_lib = Library("vllm", "FRAGMENT") # noqa
546
+
547
+
548
+def direct_register_custom_op(
549
+ op_name: str,
550
+ op_func: Callable,
551
+ mutates_args: list[str] | None = None,
552
+ fake_impl: Callable | None = None,
553
+ target_lib: Library | None = None,
554
+ dispatch_key: str | None = None,
555
+ tags: tuple[torch.Tag, ...] = (),
556
+):
557
+ """
558
+ `torch.library.custom_op` can have significant overhead because it
559
+ needs to consider complicated dispatching logic. This function
560
+ directly registers a custom op and dispatches it to the CUDA backend.
561
+ See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
562
+ for more details.
563
+
564
+ By default, the custom op is registered to the vLLM library. If you
565
+ want to register it to a different library, you can pass the library
566
+ object to the `target_lib` argument.
567
+
568
+ IMPORTANT: the lifetime of the operator is tied to the lifetime of the
569
+ library object. If you want to bind the operator to a different library,
570
+ make sure the library object is alive when the operator is used.
571
+ """
572
+ if not supports_custom_op():
573
+ from vllm.platforms import current_platform
574
+
575
+ assert not current_platform.is_cuda_alike(), (
576
+ "cuda platform needs torch>=2.4 to support custom op, "
577
+ "chances are you are using an old version of pytorch "
578
+ "or a custom build of pytorch. It is recommended to "
579
+ "use vLLM in a fresh new environment and let it install "
580
+ "the required dependencies."
581
+ )
582
+ return
583
+
584
+ if mutates_args is None:
585
+ mutates_args = []
586
+
587
+ if dispatch_key is None:
588
+ from vllm.platforms import current_platform
589
+
590
+ dispatch_key = current_platform.dispatch_key
591
+
592
+ import torch.library
593
+
594
+ if hasattr(torch.library, "infer_schema"):
595
+ schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
596
+ else:
597
+ # for pytorch 2.4
598
+ import torch._custom_op.impl
599
+
600
+ schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
601
+ my_lib = target_lib or vllm_lib
602
+ my_lib.define(op_name + schema_str, tags=tags)
603
+ my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
604
+ if fake_impl is not None:
605
+ my_lib._register_fake(op_name, fake_impl)
vllm/v1/attention/backends/flex_attention.py
+2
-1
@@ -28,7 +28,8 @@ from vllm.logger import init_logger
28
from vllm.model_executor.layers.batch_invariant import (
29
vllm_is_batch_invariant,
30
)
31
-from vllm.utils import cdiv, is_torch_equal_or_newer
31
+from vllm.utils import cdiv
32
+from vllm.utils.torch_utils import is_torch_equal_or_newer
33
from vllm.v1.attention.backends.utils import (
34
AttentionMetadataBuilder,
35
CommonAttentionMetadata,
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
-1
@@ -29,7 +29,7 @@ if current_platform.is_rocm():
29
import aiter
30
31
from vllm.triton_utils import tl, triton
32
- from vllm.utils import direct_register_custom_op
32
+ from vllm.utils.torch_utils import direct_register_custom_op
33
34
@triton.jit
35
def _vllm_layout_trans_kernel(
vllm/v1/kv_cache_interface.py
+2
-1
@@ -10,7 +10,8 @@ from typing_extensions import Self
10
11
from vllm.config import VllmConfig
12
from vllm.logger import init_logger
13
-from vllm.utils import cdiv, get_dtype_size
13
+from vllm.utils import cdiv
14
+from vllm.utils.torch_utils import get_dtype_size
15
16
logger = init_logger(__name__)
17
vllm/v1/sample/ops/penalties.py
+2
-1
@@ -4,7 +4,8 @@
4
import torch
5
6
from vllm.model_executor.layers.utils import apply_penalties
7
-from vllm.utils import is_pin_memory_available, make_tensor_with_pad
7
+from vllm.utils import is_pin_memory_available
8
+from vllm.utils.torch_utils import make_tensor_with_pad
9
10
11
def apply_all_penalties(
vllm/v1/worker/gpu_model_runner.py
+5
-3
@@ -71,16 +71,18 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
71
from vllm.utils import (
72
cdiv,
73
check_use_alibi,
74
- get_dtype_size,
74
is_pin_memory_available,
76
- kv_cache_dtype_str_to_dtype,
75
length_from_prompt_token_ids_or_embeds,
76
round_up,
79
- supports_dynamo,
77
)
78
from vllm.utils.jsontree import json_map_leaves
79
from vllm.utils.mem_constants import GiB_bytes
80
from vllm.utils.mem_utils import DeviceMemoryProfiler
81
+from vllm.utils.torch_utils import (
82
+ get_dtype_size,
83
+ kv_cache_dtype_str_to_dtype,
84
+ supports_dynamo,
85
+)
86
from vllm.v1.attention.backends.flash_attn import AttentionMetadata
87
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
88
from vllm.v1.attention.backends.utils import (
vllm/v1/worker/tpu_worker.py
+2
-1
@@ -25,7 +25,8 @@ from vllm.model_executor import set_random_seed
25
from vllm.platforms import current_platform
26
from vllm.platforms.tpu import USE_TPU_INFERENCE
27
from vllm.tasks import SupportedTask
28
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
28
+from vllm.utils import cdiv
29
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
30
from vllm.v1.core.sched.output import SchedulerOutput
31
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec
32
from vllm.v1.outputs import ModelRunnerOutput
vllm/v1/worker/ubatching.py
+1
-1
@@ -7,7 +7,7 @@ import torch
7
8
from vllm import forward_context
9
from vllm.forward_context import ForwardContext
10
-from vllm.utils import current_stream
10
+from vllm.utils.torch_utils import current_stream
11
12
_THREAD_ID_TO_CONTEXT: dict = {}
13
_CURRENT_CONTEXTS: list[Optional["UBatchContext"]] = [None, None]