Skip to content

Commit 56a1ad4

Browse files
committed
Update new chunked prefill distributed test to include non-Ray
1 parent 1938c35 commit 56a1ad4

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

tests/distributed/test_chunked_prefill_distributed.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
@pytest.mark.parametrize("dtype", ["half"])
2828
@pytest.mark.parametrize("max_tokens", [5])
2929
@pytest.mark.parametrize("chunked_prefill_token_size", [16])
30+
@pytest.mark.parametrize("worker_use_ray", [False, True])
3031
def test_models(
3132
hf_runner,
3233
vllm_runner,
@@ -35,6 +36,7 @@ def test_models(
3536
dtype: str,
3637
max_tokens: int,
3738
chunked_prefill_token_size: int,
39+
worker_use_ray: bool,
3840
) -> None:
3941
# Add a chunked prefill config.
4042
max_num_seqs = min(chunked_prefill_token_size, 256)
@@ -53,6 +55,7 @@ def test_models(
5355
max_num_seqs=max_num_seqs,
5456
enable_chunked_prefill=enable_chunked_prefill,
5557
max_num_batched_tokens=max_num_batched_tokens,
58+
worker_use_ray=worker_use_ray,
5659
)
5760
vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
5861
del vllm_model

0 commit comments

Comments
 (0)