ca76486a16 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/core/test_uva.py

+1 -1

@@ -3,7 +3,7 @@

3 import pytest

4 import torch

5

6 -from vllm.utils import is_uva_available

6 +from vllm.utils.platform_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)]

tests/v1/logits_processors/test_correctness.py

+1 -1

@@ -21,7 +21,7 @@ from tests.v1.sample.utils import (

21 from vllm.config import VllmConfig

22 from vllm.platforms import current_platform

23 from vllm.sampling_params import SamplingParams

24 -from vllm.utils import is_pin_memory_available

24 +from vllm.utils.platform_utils import is_pin_memory_available

25 from vllm.v1.sample.logits_processor import (

26 BatchUpdate,

27 BatchUpdateBuilder,

tests/v1/sample/test_sampler.py

+1 -1

@@ -7,7 +7,7 @@ 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

10 +from vllm.utils.platform_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

tests/v1/worker/test_gpu_input_batch.py

+1 -1

@@ -10,7 +10,7 @@ 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

13 +from vllm.utils.platform_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

vllm/device_allocator/cumem.py

+1 -1

@@ -18,7 +18,7 @@ from typing import Any

18 import torch

19

20 from vllm.logger import init_logger

21 -from vllm.utils import is_pin_memory_available

21 +from vllm.utils.platform_utils import is_pin_memory_available

22

23 logger = init_logger(__name__)

24

vllm/lora/lora_weights.py

+1 -1

@@ -8,7 +8,7 @@ import torch

8 import torch.types

9

10 from vllm.lora.peft_helper import PEFTHelper

11 -from vllm.utils import is_pin_memory_available

11 +from vllm.utils.platform_utils import is_pin_memory_available

12

13

14 class LoRALayerWeights:

vllm/lora/models.py

+1 -1

@@ -31,8 +31,8 @@ from vllm.model_executor.models import SupportsLoRA, supports_multimodal

31 from vllm.model_executor.models.interfaces import is_pooling_model

32 from vllm.model_executor.models.module_mapping import MultiModelKeys

33 from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper

34 -from vllm.utils import is_pin_memory_available

34 from vllm.utils.cache import LRUCache

35 +from vllm.utils.platform_utils import is_pin_memory_available

36

37 logger = init_logger(__name__)

38

vllm/model_executor/model_loader/utils.py

+1 -1

@@ -26,7 +26,7 @@ from vllm.model_executor.models.adapters import (

26 try_create_mm_pooling_model_cls,

27 )

28 from vllm.model_executor.models.interfaces import SupportsQuant, supports_multimodal

29 -from vllm.utils import is_pin_memory_available

29 +from vllm.utils.platform_utils import is_pin_memory_available

30

31 logger = init_logger(__name__)

32

vllm/model_executor/models/qwen2_5_vl.py

+1 -1

@@ -72,7 +72,7 @@ from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargs

72 from vllm.multimodal.parse import MultiModalDataItems

73 from vllm.multimodal.processing import PromptReplacement, PromptUpdate

74 from vllm.sequence import IntermediateTensors

75 -from vllm.utils import is_pin_memory_available

75 +from vllm.utils.platform_utils import is_pin_memory_available

76 from vllm.utils.tensor_schema import TensorSchema, TensorShape

77

78 from .interfaces import (

vllm/model_executor/models/utils.py

+2 -2

@@ -22,8 +22,8 @@ from vllm.logger import init_logger

22 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

23 from vllm.multimodal import NestedTensors

24 from vllm.sequence import IntermediateTensors

25 -from vllm.utils import (

26 - cdiv,

25 +from vllm.utils import cdiv

26 +from vllm.utils.platform_utils import (

27 is_pin_memory_available,

28 is_uva_available,

29 )

vllm/usage/usage_lib.py

+1 -1

@@ -21,7 +21,7 @@ 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_get_device_properties

24 +from vllm.utils.platform_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

vllm/utils/__init__.py

+3 -50

@@ -27,12 +27,8 @@ from argparse import (

27 _ArgumentGroup,

28 )

29 from collections import defaultdict

30 -from collections.abc import (

31 - Callable,

32 - Sequence,

33 -)

34 -from concurrent.futures.process import ProcessPoolExecutor

35 -from functools import cache, partial, wraps

30 +from collections.abc import Callable

31 +from functools import partial, wraps

32 from typing import TYPE_CHECKING, Any, TypeVar

33

34 import cloudpickle

@@ -44,6 +40,7 @@ import yaml

40 import vllm.envs as envs

41 from vllm.logger import enable_trace_function_call, init_logger

42 from vllm.ray.lazy_utils import is_in_ray_actor

43 +from vllm.utils.platform_utils import cuda_is_initialized, xpu_is_initialized

44

45 _DEPRECATED_MAPPINGS = {

46 "cprofile": "profiling",

@@ -169,21 +166,6 @@ def round_down(x: int, y: int) -> int:

166 return (x // y) * y

167

168

172 -@cache

173 -def is_pin_memory_available() -> bool:

174 - from vllm.platforms import current_platform

175 -

176 - return current_platform.is_pin_memory_available()

177 -

178 -

179 -@cache

180 -def is_uva_available() -> bool:

181 - """Check if Unified Virtual Addressing (UVA) is available."""

182 - # UVA requires pinned memory.

183 - # TODO: Add more requirements for UVA if needed.

184 - return is_pin_memory_available()

185 -

186 -

169 # TODO: This function can be removed if transformer_modules classes are

170 # serialized by value when communicating between processes

171 def init_cached_hf_modules() -> None:

@@ -216,35 +198,6 @@ def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:

198 enable_trace_function_call(log_path)

199

200

219 -def cuda_is_initialized() -> bool:

220 - """Check if CUDA is initialized."""

221 - if not torch.cuda._is_compiled():

222 - return False

223 - return torch.cuda.is_initialized()

224 -

225 -

226 -def xpu_is_initialized() -> bool:

227 - """Check if XPU is initialized."""

228 - if not torch.xpu._is_compiled():

229 - return False

230 - return torch.xpu.is_initialized()

231 -

232 -

233 -def cuda_get_device_properties(

234 - device, names: Sequence[str], init_cuda=False

235 -) -> tuple[Any, ...]:

236 - """Get specified CUDA device property values without initializing CUDA in

237 - the current process."""

238 - if init_cuda or cuda_is_initialized():

239 - props = torch.cuda.get_device_properties(device)

240 - return tuple(getattr(props, name) for name in names)

241 -

242 - # Run in subprocess to avoid initializing CUDA as a side effect.

243 - mp_ctx = multiprocessing.get_context("fork")

244 - with ProcessPoolExecutor(max_workers=1, mp_context=mp_ctx) as executor:

245 - return executor.submit(cuda_get_device_properties, device, names, True).result()

246 -

247 -

201 def weak_bind(

202 bound_method: Callable[..., Any],

203 ) -> Callable[..., None]:

vllm/utils/platform_utils.py new

+54

@@ -0,0 +1,54 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import multiprocessing

5 +from collections.abc import Sequence

6 +from concurrent.futures.process import ProcessPoolExecutor

7 +from functools import cache

8 +from typing import Any

9 +

10 +import torch

11 +

12 +

13 +def cuda_is_initialized() -> bool:

14 + """Check if CUDA is initialized."""

15 + if not torch.cuda._is_compiled():

16 + return False

17 + return torch.cuda.is_initialized()

18 +

19 +

20 +def xpu_is_initialized() -> bool:

21 + """Check if XPU is initialized."""

22 + if not torch.xpu._is_compiled():

23 + return False

24 + return torch.xpu.is_initialized()

25 +

26 +

27 +def cuda_get_device_properties(

28 + device, names: Sequence[str], init_cuda=False

29 +) -> tuple[Any, ...]:

30 + """Get specified CUDA device property values without initializing CUDA in

31 + the current process."""

32 + if init_cuda or cuda_is_initialized():

33 + props = torch.cuda.get_device_properties(device)

34 + return tuple(getattr(props, name) for name in names)

35 +

36 + # Run in subprocess to avoid initializing CUDA as a side effect.

37 + mp_ctx = multiprocessing.get_context("fork")

38 + with ProcessPoolExecutor(max_workers=1, mp_context=mp_ctx) as executor:

39 + return executor.submit(cuda_get_device_properties, device, names, True).result()

40 +

41 +

42 +@cache

43 +def is_pin_memory_available() -> bool:

44 + from vllm.platforms import current_platform

45 +

46 + return current_platform.is_pin_memory_available()

47 +

48 +

49 +@cache

50 +def is_uva_available() -> bool:

51 + """Check if Unified Virtual Addressing (UVA) is available."""

52 + # UVA requires pinned memory.

53 + # TODO: Add more requirements for UVA if needed.

54 + return is_pin_memory_available()

vllm/v1/attention/backends/flashinfer.py

+2 -1

@@ -34,12 +34,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

34 )

35 from vllm.platforms import current_platform

36 from vllm.triton_utils import tl, triton

37 -from vllm.utils import cdiv, is_pin_memory_available

37 +from vllm.utils import cdiv

38 from vllm.utils.flashinfer import (

39 can_use_trtllm_attention,

40 flashinfer_disable_q_quantization,

41 use_trtllm_attention,

42 )

43 +from vllm.utils.platform_utils import is_pin_memory_available

44 from vllm.v1.attention.backends.utils import (

45 AttentionCGSupport,

46 AttentionMetadataBuilder,

vllm/v1/kv_offload/worker/cpu_gpu.py

+1 -1

@@ -7,7 +7,7 @@ import torch

7 from vllm import _custom_ops as ops

8 from vllm.attention import AttentionBackend

9 from vllm.logger import init_logger

10 -from vllm.utils import is_pin_memory_available

10 +from vllm.utils.platform_utils import is_pin_memory_available

11 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec

12 from vllm.v1.kv_offload.worker.worker import (

13 OffloadingHandler,

vllm/v1/pool/metadata.py

+1 -1

@@ -5,7 +5,7 @@ from dataclasses import dataclass

5 import torch

6

7 from vllm.pooling_params import PoolingParams

8 -from vllm.utils import is_pin_memory_available

8 +from vllm.utils.platform_utils import is_pin_memory_available

9

10 pin_memory = is_pin_memory_available()

11

vllm/v1/sample/ops/penalties.py

+1 -1

@@ -4,7 +4,7 @@

4 import torch

5

6 from vllm.model_executor.layers.utils import apply_penalties

7 -from vllm.utils import is_pin_memory_available

7 +from vllm.utils.platform_utils import is_pin_memory_available

8 from vllm.utils.torch_utils import make_tensor_with_pad

9

10

vllm/v1/sample/sampler.py

+1 -1

@@ -6,7 +6,7 @@ import torch

6 import torch.nn as nn

7

8 from vllm.config.model import LogprobsMode

9 -from vllm.utils import is_pin_memory_available

9 +from vllm.utils.platform_utils import is_pin_memory_available

10 from vllm.v1.outputs import LogprobsTensors, SamplerOutput

11 from vllm.v1.sample.metadata import SamplingMetadata

12 from vllm.v1.sample.ops.bad_words import apply_bad_words

vllm/v1/spec_decode/eagle.py

+1 -1

@@ -24,7 +24,7 @@ from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache

24 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

25 from vllm.multimodal import MULTIMODAL_REGISTRY

26 from vllm.platforms import current_platform

27 -from vllm.utils import is_pin_memory_available

27 +from vllm.utils.platform_utils import is_pin_memory_available

28 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

29 from vllm.v1.attention.backends.tree_attn import (

30 TreeAttentionMetadata,

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -72,13 +72,13 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

72 from vllm.utils import (

73 cdiv,

74 check_use_alibi,

75 - is_pin_memory_available,

75 length_from_prompt_token_ids_or_embeds,

76 round_up,

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.platform_utils import is_pin_memory_available

82 from vllm.utils.torch_utils import (

83 get_dtype_size,

84 kv_cache_dtype_str_to_dtype,

vllm/v1/worker/tpu_model_runner.py

+2 -1

@@ -53,7 +53,8 @@ from vllm.multimodal.inputs import (

53 from vllm.multimodal.utils import group_mm_kwargs_by_modality

54 from vllm.sequence import IntermediateTensors

55 from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

56 -from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available, prev_power_of_2

56 +from vllm.utils import LayerBlockType, cdiv, prev_power_of_2

57 +from vllm.utils.platform_utils import is_pin_memory_available

58 from vllm.v1.attention.backends.pallas import (

59 TPU_STR_DTYPE_TO_TORCH_DTYPE,

60 PallasAttentionBackend,

[Chore] Separate out `vllm.utils.platform_utils.py` (#27374)