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
241ad7b3
Unverified
Commit
241ad7b3
authored
Jan 10, 2025
by
youkaichao
Committed by
GitHub
Jan 10, 2025
Browse files
[ci] Fix sampler tests (#11922)
Signed-off-by:
youkaichao
<
youkaichao@gmail.com
>
parent
d85c47d6
Changes
2
Show whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
10 additions
and
2 deletions
+10
-2
.buildkite/test-pipeline.yaml
.buildkite/test-pipeline.yaml
+1
-0
tests/conftest.py
tests/conftest.py
+9
-2
No files found.
.buildkite/test-pipeline.yaml
View file @
241ad7b3
...
...
@@ -214,6 +214,7 @@ steps:
-
vllm/model_executor/layers
-
vllm/sampling_metadata.py
-
tests/samplers
-
tests/conftest.py
commands
:
-
pytest -v -s samplers
-
VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
...
...
tests/conftest.py
View file @
241ad7b3
...
...
@@ -28,12 +28,13 @@ from vllm.distributed import (cleanup_dist_env_and_memory,
init_distributed_environment
,
initialize_model_parallel
)
from
vllm.inputs
import
(
ExplicitEncoderDecoderPrompt
,
TextPrompt
,
to_enc_dec_tuple_list
,
zip_enc_dec_prompts
)
TokensPrompt
,
to_enc_dec_tuple_list
,
zip_enc_dec_prompts
)
from
vllm.logger
import
init_logger
from
vllm.outputs
import
RequestOutput
from
vllm.sampling_params
import
BeamSearchParams
from
vllm.utils
import
(
STR_DTYPE_TO_TORCH_DTYPE
,
cuda_device_count_stateless
,
identity
)
identity
,
is_list_of
)
logger
=
init_logger
(
__name__
)
...
...
@@ -886,6 +887,12 @@ class VllmRunner:
beam_width
:
int
,
max_tokens
:
int
,
)
->
List
[
Tuple
[
List
[
List
[
int
]],
List
[
str
]]]:
if
is_list_of
(
prompts
,
str
,
check
=
"all"
):
prompts
=
[
TextPrompt
(
prompt
=
prompt
)
for
prompt
in
prompts
]
else
:
prompts
=
[
TokensPrompt
(
prompt_token_ids
=
tokens
)
for
tokens
in
prompts
]
outputs
=
self
.
model
.
beam_search
(
prompts
,
BeamSearchParams
(
beam_width
=
beam_width
,
max_tokens
=
max_tokens
))
...
...
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