File tree 3 files changed +4
-8
lines changed
3 files changed +4
-8
lines changed Original file line number Diff line number Diff line change @@ -287,8 +287,7 @@ class HfRunner:
287
287
def get_default_device (self ):
288
288
from vllm .platforms import current_platform
289
289
290
- return ("cpu"
291
- if current_platform .is_cpu () else current_platform .device_type )
290
+ return ("cpu" if current_platform .is_cpu () else "cuda" )
292
291
293
292
def wrap_device (self , x : _T , device : Optional [str ] = None ) -> _T :
294
293
if x is None or isinstance (x , (bool , )):
Original file line number Diff line number Diff line change 6
6
import pytest
7
7
import torch
8
8
9
- from vllm .platforms import current_platform
10
9
from vllm .utils import make_tensor_with_pad
11
10
from vllm .v1 .sample .metadata import SamplingMetadata
12
11
from vllm .v1 .sample .sampler import Sampler
13
12
14
13
VOCAB_SIZE = 1024
15
14
NUM_OUTPUT_TOKENS = 20
16
15
CUDA_DEVICES = [
17
- f"{ current_platform .device_type } :{ i } "
18
- for i in range (1 if current_platform .device_count () == 1 else 2 )
16
+ f"cuda:{ i } " for i in range (1 if torch .cuda .device_count () == 1 else 2 )
19
17
]
20
18
MAX_NUM_PROMPT_TOKENS = 64
21
19
Original file line number Diff line number Diff line change 14
14
SamplerOutput ,
15
15
SamplingMetadata , get_logprobs ,
16
16
get_pythonized_sample_results )
17
- from vllm .platforms import current_platform
18
17
from vllm .sequence import (CompletionSequenceGroupOutput , IntermediateTensors ,
19
18
Logprob , SequenceGroupMetadata , SequenceOutput )
20
19
from vllm .utils import PyObjectCache , async_tensor_h2d , current_stream
@@ -159,8 +158,8 @@ class StatefulModelInput(BroadcastableModelInput):
159
158
is_first_multi_step : bool = False
160
159
base_output_proc_callback : Optional [Callable ] = None
161
160
# ping-pong data structures for multi-step to wait on the previous step
162
- step_cuda_events : List [current_platform .Event ] = field (
163
- default_factory = lambda : [current_platform .Event (blocking = True )] * 2 )
161
+ step_cuda_events : List [torch . cuda .Event ] = field (
162
+ default_factory = lambda : [torch . cuda .Event (blocking = True )] * 2 )
164
163
num_seqs : int = - 1
165
164
num_queries : int = - 1
166
165
num_single_step_prefills : int = 0
You can’t perform that action at this time.
0 commit comments