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 6, 05:45 UTC
[ROCm][CI] Add large_gpu_mark to test_max_tokens_none for ROCm (#37717)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 21, 2026 at 23:25 UTC
c86b17cfe663afb46a46b8081e9fd4c81dc920a2
1 file changed
+14
-2
tests/test_regression.py
+14
-2
@@ -12,6 +12,7 @@ import gc
12
import pytest
13
import torch
14
15
+from tests.utils import large_gpu_mark
16
from vllm import LLM, SamplingParams
17
from vllm.platforms import current_platform
18
@@ -32,10 +33,21 @@ def test_duplicated_ignored_sequence_group():
33
assert len(prompts) == len(outputs)
34
35
35
-def test_max_tokens_none():
36
+@pytest.mark.parametrize(
37
+ "model",
38
+ [
39
+ pytest.param(
40
+ "distilbert/distilgpt2",
41
+ marks=[
42
+ *([large_gpu_mark(min_gb=80)] if current_platform.is_rocm() else []),
43
+ ],
44
+ ),
45
+ ],
46
+)
47
+def test_max_tokens_none(model):
48
sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=None)
49
llm = LLM(
38
- model="distilbert/distilgpt2",
50
+ model=model,
51
max_num_batched_tokens=4096,
52
tensor_parallel_size=1,
53
)