Skip to content

Commit 3f1fc74

Browse files
authored
[V1][CI/Test] Do basic test for top-p & top-k sampling (vllm-project#12469)
Signed-off-by: Woosuk Kwon <[email protected]>
1 parent 01ba927 commit 3f1fc74

File tree

1 file changed

+20
-8
lines changed

1 file changed

+20
-8
lines changed

tests/v1/engine/test_engine_core.py

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@ def test_engine_core(monkeypatch):
144144
def test_engine_core_advanced_sampling(monkeypatch):
145145
"""
146146
A basic end-to-end test to verify that the engine functions correctly
147-
when additional sampling parameters, such as min_tokens and
147+
when additional sampling parameters, such as top_p, min_tokens, and
148148
presence_penalty, are set.
149149
"""
150150
with monkeypatch.context() as m:
@@ -167,11 +167,23 @@ def test_engine_core_advanced_sampling(monkeypatch):
167167
stop_token_ids=[1001, 1002],
168168
)
169169
engine_core.add_request(request)
170-
assert len(engine_core.scheduler.waiting) == 1
171-
assert len(engine_core.scheduler.running) == 0
172-
# Loop through until they are all done.
173-
while len(engine_core.step().outputs) > 0:
174-
pass
175170

176-
assert len(engine_core.scheduler.waiting) == 0
177-
assert len(engine_core.scheduler.running) == 0
171+
def _check_engine_state():
172+
assert len(engine_core.scheduler.waiting) == 1
173+
assert len(engine_core.scheduler.running) == 0
174+
# Loop through until they are all done.
175+
while len(engine_core.step().outputs) > 0:
176+
pass
177+
assert len(engine_core.scheduler.waiting) == 0
178+
assert len(engine_core.scheduler.running) == 0
179+
180+
_check_engine_state()
181+
182+
# Second request.
183+
request2 = make_request()
184+
request2.sampling_params = SamplingParams(
185+
top_p=0.99,
186+
top_k=50,
187+
)
188+
engine_core.add_request(request2)
189+
_check_engine_state()

0 commit comments

Comments
 (0)