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 15, 03:30 UTC
[UT][intel GPU] use current_platform instead of device hardcode in v1 tests (#20169)
Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com>
Liangliang Ma committed
Jul 2, 2025 at 09:06 UTC
a0389e0554728b8478503de9a1f350e514635fe1
10 files changed
+44
-26
tests/conftest.py
+2
-2
@@ -34,7 +34,6 @@ from vllm.logger import init_logger
34
from vllm.outputs import RequestOutput
35
from vllm.sampling_params import BeamSearchParams
36
from vllm.transformers_utils.utils import maybe_model_redirect
37
-from vllm.utils import cuda_device_count_stateless
37
38
logger = init_logger(__name__)
39
@@ -1094,7 +1093,8 @@ def num_gpus_available():
1093
"""Get number of GPUs without initializing the CUDA context
1094
in current process."""
1095
1097
- return cuda_device_count_stateless()
1096
+ from vllm.platforms import current_platform
1097
+ return current_platform.device_count()
1098
1099
1100
temp_dir = tempfile.gettempdir()
tests/v1/sample/test_rejection_sampler.py
+5
-4
@@ -6,12 +6,13 @@ import pytest
6
import torch
7
import torch.nn.functional as F
8
9
+from vllm.platforms import current_platform
10
from vllm.v1.sample.metadata import SamplingMetadata
11
from vllm.v1.sample.rejection_sampler import (PLACEHOLDER_TOKEN_ID,
12
RejectionSampler)
13
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
14
14
-DEVICE = "cuda"
15
+DEVICE = current_platform.device_type
16
17
18
@pytest.fixture
@@ -21,7 +22,7 @@ def rejection_sampler():
22
23
def create_logits_tensor(output_token_ids: list[list[int]],
24
vocab_size: int = 100) -> torch.Tensor:
24
- """Helper function to create logits tensor that
25
+ """Helper function to create logits tensor that
26
will produce desired token ids on argmax"""
27
token_ids = [tokens[:-1] for tokens in output_token_ids]
28
num_total_tokens = sum(len(tokens) for tokens in token_ids)
@@ -41,8 +42,8 @@ def create_sampling_metadata(
42
top_p: Optional[torch.Tensor] = None,
43
generators: Optional[dict[int, Any]] = None,
44
) -> SamplingMetadata:
44
- """Create a v1 sampling metadata object with all_greedy set
45
- to the given value. Either all greedy or all random sampling
45
+ """Create a v1 sampling metadata object with all_greedy set
46
+ to the given value. Either all greedy or all random sampling
47
is used.
48
"""
49
generators = generators or {}
tests/v1/sample/test_topk_topp_sampler.py
+6
-5
@@ -2,25 +2,26 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
import torch
5
-from flashinfer.sampling import top_k_renorm_probs, top_p_renorm_probs
5
from torch import Generator
6
7
from vllm.platforms import current_platform
8
from vllm.v1.sample.ops.topk_topp_sampler import (apply_top_k_top_p,
9
is_flashinfer_available)
10
12
-DEVICE = "cuda"
11
+DEVICE = current_platform.device_type
12
13
BATCH_SIZE = 1024
14
VOCAB_SIZE = 128 * 1024
15
16
FLASHINFER_ENABLED = current_platform.is_cuda() and is_flashinfer_available
17
+if is_flashinfer_available:
18
+ from flashinfer.sampling import top_k_renorm_probs, top_p_renorm_probs
19
20
21
@pytest.fixture(autouse=True)
22
def reset_default_device():
23
"""
23
- Explicitly set the default device, which can affect subsequent tests.
24
+ Explicitly set the default device, which can affect subsequent tests.
25
Adding this fixture helps avoid this problem.
26
"""
27
original_device = torch.get_default_device()
@@ -58,8 +59,8 @@ def test_flashinfer_sampler():
59
This test verifies that the FlashInfer top-k and top-p sampling
60
implementation produces the same results as the Python implementation.
61
61
- NOTE: FlashInfer did not directly expose an interface for fused top-k and
62
- top-p prob renorm (it did provide fused sampling but we cannot compare
62
+ NOTE: FlashInfer did not directly expose an interface for fused top-k and
63
+ top-p prob renorm (it did provide fused sampling but we cannot compare
64
sampling results due to randomness), so we will compare the probability
65
renormed consequently by top-k and then top-p of FlashInfer implementation.
66
'''
tests/v1/spec_decode/test_eagle.py
+13
-10
@@ -10,6 +10,7 @@ from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, ModelConfig,
10
ParallelConfig, SchedulerConfig, SpeculativeConfig,
11
VllmConfig)
12
from vllm.model_executor.models.llama import LlamaForCausalLM
13
+from vllm.platforms import current_platform
14
from vllm.v1.spec_decode.eagle import EagleProposer
15
16
model_dir = "meta-llama/Llama-3.1-8B-Instruct"
@@ -38,15 +39,17 @@ def _create_proposer(method: str, k: int) -> EagleProposer:
39
num_speculative_tokens=k,
40
)
41
41
- vllm_config = VllmConfig(model_config=model_config,
42
- cache_config=CacheConfig(),
43
- speculative_config=speculative_config,
44
- device_config=DeviceConfig(device="cuda"),
45
- parallel_config=ParallelConfig(),
46
- load_config=LoadConfig(),
47
- scheduler_config=SchedulerConfig())
42
+ vllm_config = VllmConfig(
43
+ model_config=model_config,
44
+ cache_config=CacheConfig(),
45
+ speculative_config=speculative_config,
46
+ device_config=DeviceConfig(device=current_platform.device_type),
47
+ parallel_config=ParallelConfig(),
48
+ load_config=LoadConfig(),
49
+ scheduler_config=SchedulerConfig())
50
49
- return EagleProposer(vllm_config=vllm_config, device='cuda')
51
+ return EagleProposer(vllm_config=vllm_config,
52
+ device=current_platform.device_type)
53
54
55
def test_prepare_inputs():
@@ -59,7 +62,7 @@ def test_prepare_inputs():
62
a, a + 1, ..., a + b - n2 - 1,
63
a + b, a + b + 1, ..., a + b + c - n3 - 1]
64
"""
62
- device = torch.device('cuda')
65
+ device = torch.device(current_platform.device_type)
66
67
# a = 4, b = 7, c = 5
68
# n1 = 1, n2 = 3, n3 = 2
@@ -198,7 +201,7 @@ def test_load_model(mock_get_model, mock_get_layers, mock_get_pp_group, method,
201
@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 8])
202
def test_propose(num_speculative_tokens):
203
# Use GPU device
201
- device = torch.device('cuda')
204
+ device = torch.device(current_platform.device_type)
205
206
# Setup test parameters
207
batch_size = 2
tests/v1/worker/test_gpu_input_batch.py
+3
-1
@@ -8,6 +8,7 @@ import numpy as np
8
import pytest
9
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
14
from vllm.v1.pool.metadata import PoolingMetadata
@@ -19,7 +20,8 @@ VOCAB_SIZE = 1024
20
NUM_OUTPUT_TOKENS = 20
21
MAX_PROMPT_SIZE = 100
22
CUDA_DEVICES = [
22
- f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
23
+ f"{current_platform.device_type}:{i}"
24
+ for i in range(min(current_platform.device_count(), 2))
25
]
26
MAX_NUM_PROMPT_TOKENS = 64
27
tests/v1/worker/test_gpu_model_runner.py
+2
-1
@@ -9,6 +9,7 @@ import torch
9
from vllm.attention import Attention
10
from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
11
SchedulerConfig, VllmConfig, set_current_vllm_config)
12
+from vllm.platforms import current_platform
13
from vllm.sampling_params import SamplingParams
14
from vllm.utils import GiB_bytes
15
from vllm.v1.core.kv_cache_utils import (estimate_max_model_len,
@@ -23,7 +24,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
24
25
BLOCK_SIZE = 16
26
NUM_BLOCKS = 10
26
-DEVICE = "cuda"
27
+DEVICE = current_platform.device_type
28
29
30
def initialize_kv_cache(runner: GPUModelRunner):
vllm/platforms/cuda.py
+5
-1
@@ -18,7 +18,7 @@ from typing_extensions import ParamSpec
18
import vllm._C # noqa
19
import vllm.envs as envs
20
from vllm.logger import init_logger
21
-from vllm.utils import import_pynvml
21
+from vllm.utils import cuda_device_count_stateless, import_pynvml
22
23
from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
24
@@ -401,6 +401,10 @@ class CudaPlatformBase(Platform):
401
pg._register_backend(device, backend_type, backend_class)
402
return pg
403
404
+ @classmethod
405
+ def device_count(cls) -> int:
406
+ return cuda_device_count_stateless()
407
+
408
409
# NVML utils
410
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/rocm.py
+5
@@ -12,6 +12,7 @@ from torch.distributed.distributed_c10d import is_nccl_available
12
13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
+from vllm.utils import cuda_device_count_stateless
16
17
from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
18
@@ -446,3 +447,7 @@ class RocmPlatform(Platform):
447
448
pg._register_backend(device, backend_type, backend_class)
449
return pg
450
+
451
+ @classmethod
452
+ def device_count(cls) -> int:
453
+ return cuda_device_count_stateless()
vllm/platforms/xpu.py
+1
-1
@@ -189,4 +189,4 @@ class XPUPlatform(Platform):
189
190
@classmethod
191
def device_count(cls) -> int:
192
- return torch.xpu.device_count()
192
+ return torch.xpu.device_count()
\ No newline at end of file
vllm/v1/attention/backends/mla/common.py
+2
-1
@@ -217,7 +217,8 @@ try:
217
is_vllm_fa = True
218
except ImportError:
219
# For rocm use upstream flash attention
220
- from flash_attn import flash_attn_varlen_func
220
+ if current_platform.is_rocm():
221
+ from flash_attn import flash_attn_varlen_func
222
is_vllm_fa = False
223
224
if TYPE_CHECKING: