Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
bdb2cdda
Unverified
Commit
bdb2cdda
authored
Apr 29, 2025
by
ponix-j
Committed by
GitHub
Apr 29, 2025
Browse files
[Misc]Use a platform independent interface to obtain the device attributes (#17100)
parent
ebb3930d
Changes
3
Show whitespace changes
Inline
Side-by-side
Showing
3 changed files
with
8 additions
and
4 deletions
+8
-4
tests/conftest.py
tests/conftest.py
+2
-1
tests/v1/sample/test_sampler.py
tests/v1/sample/test_sampler.py
+3
-1
vllm/worker/multi_step_model_runner.py
vllm/worker/multi_step_model_runner.py
+3
-2
No files found.
tests/conftest.py
View file @
bdb2cdda
...
@@ -293,7 +293,8 @@ class HfRunner:
...
@@ -293,7 +293,8 @@ class HfRunner:
def
get_default_device
(
self
):
def
get_default_device
(
self
):
from
vllm.platforms
import
current_platform
from
vllm.platforms
import
current_platform
return
(
"cpu"
if
current_platform
.
is_cpu
()
else
"cuda"
)
return
(
"cpu"
if
current_platform
.
is_cpu
()
else
current_platform
.
device_type
)
def
wrap_device
(
self
,
x
:
_T
,
device
:
Optional
[
str
]
=
None
)
->
_T
:
def
wrap_device
(
self
,
x
:
_T
,
device
:
Optional
[
str
]
=
None
)
->
_T
:
if
x
is
None
or
isinstance
(
x
,
(
bool
,
)):
if
x
is
None
or
isinstance
(
x
,
(
bool
,
)):
...
...
tests/v1/sample/test_sampler.py
View file @
bdb2cdda
...
@@ -6,6 +6,7 @@ import numpy as np
...
@@ -6,6 +6,7 @@ import numpy as np
import
pytest
import
pytest
import
torch
import
torch
from
vllm.platforms
import
current_platform
from
vllm.utils
import
make_tensor_with_pad
from
vllm.utils
import
make_tensor_with_pad
from
vllm.v1.sample.metadata
import
SamplingMetadata
from
vllm.v1.sample.metadata
import
SamplingMetadata
from
vllm.v1.sample.sampler
import
Sampler
from
vllm.v1.sample.sampler
import
Sampler
...
@@ -13,7 +14,8 @@ from vllm.v1.sample.sampler import Sampler
...
@@ -13,7 +14,8 @@ from vllm.v1.sample.sampler import Sampler
VOCAB_SIZE
=
1024
VOCAB_SIZE
=
1024
NUM_OUTPUT_TOKENS
=
20
NUM_OUTPUT_TOKENS
=
20
CUDA_DEVICES
=
[
CUDA_DEVICES
=
[
f
"cuda:
{
i
}
"
for
i
in
range
(
1
if
torch
.
cuda
.
device_count
()
==
1
else
2
)
f
"
{
current_platform
.
device_type
}
:
{
i
}
"
for
i
in
range
(
1
if
current_platform
.
device_count
()
==
1
else
2
)
]
]
MAX_NUM_PROMPT_TOKENS
=
64
MAX_NUM_PROMPT_TOKENS
=
64
...
...
vllm/worker/multi_step_model_runner.py
View file @
bdb2cdda
...
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.sampler import (PromptLogprobs, SampleLogprobs,
...
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.sampler import (PromptLogprobs, SampleLogprobs,
SamplerOutput
,
SamplerOutput
,
SamplingMetadata
,
get_logprobs
,
SamplingMetadata
,
get_logprobs
,
get_pythonized_sample_results
)
get_pythonized_sample_results
)
from
vllm.platforms
import
current_platform
from
vllm.sequence
import
(
CompletionSequenceGroupOutput
,
IntermediateTensors
,
from
vllm.sequence
import
(
CompletionSequenceGroupOutput
,
IntermediateTensors
,
Logprob
,
SequenceGroupMetadata
,
SequenceOutput
)
Logprob
,
SequenceGroupMetadata
,
SequenceOutput
)
from
vllm.utils
import
PyObjectCache
,
async_tensor_h2d
,
current_stream
from
vllm.utils
import
PyObjectCache
,
async_tensor_h2d
,
current_stream
...
@@ -158,8 +159,8 @@ class StatefulModelInput(BroadcastableModelInput):
...
@@ -158,8 +159,8 @@ class StatefulModelInput(BroadcastableModelInput):
is_first_multi_step
:
bool
=
False
is_first_multi_step
:
bool
=
False
base_output_proc_callback
:
Optional
[
Callable
]
=
None
base_output_proc_callback
:
Optional
[
Callable
]
=
None
# ping-pong data structures for multi-step to wait on the previous step
# ping-pong data structures for multi-step to wait on the previous step
step_cuda_events
:
List
[
torch
.
cuda
.
Event
]
=
field
(
step_cuda_events
:
List
[
current_platform
.
Event
]
=
field
(
default_factory
=
lambda
:
[
torch
.
cuda
.
Event
(
blocking
=
True
)]
*
2
)
default_factory
=
lambda
:
[
current_platform
.
Event
(
blocking
=
True
)]
*
2
)
num_seqs
:
int
=
-
1
num_seqs
:
int
=
-
1
num_queries
:
int
=
-
1
num_queries
:
int
=
-
1
num_single_step_prefills
:
int
=
0
num_single_step_prefills
:
int
=
0
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment