@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 167 lines 5.45 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm import CompletionOutput, RequestOutput
7 from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin
8 from vllm.inputs import TokensInput
9 from vllm.logprobs import Logprob, SampleLogprobs
10 from vllm.sampling_params import BeamSearchParams
11
12
13 class _Tokenizer:
14 eos_token_id = 0
15 special_token_id = eos_token_id
16
17 def decode(self, token_ids: list[int], skip_special_tokens: bool = False) -> str:
18 if skip_special_tokens:
19 token_ids = [
20 token_id for token_id in token_ids if token_id != self.special_token_id
21 ]
22 return " ".join(str(token_id) for token_id in token_ids)
23
24
25 class _Renderer:
26 def get_tokenizer(self) -> _Tokenizer:
27 return _Tokenizer()
28
29
30 class _EngineClient:
31 async def generate(self, prompt, *args, **kwargs):
32 yield RequestOutput(
33 request_id=kwargs.get("request_id", "test-request"),
34 prompt=prompt.get("prompt"),
35 prompt_token_ids=prompt["prompt_token_ids"],
36 prompt_logprobs=None,
37 outputs=[
38 CompletionOutput(
39 index=0,
40 text="",
41 token_ids=[],
42 cumulative_logprob=None,
43 logprobs=[
44 {
45 11: Logprob(logprob=-1.0),
46 12: Logprob(logprob=-2.0),
47 13: Logprob(logprob=-3.0),
48 14: Logprob(logprob=-4.0),
49 _Tokenizer.eos_token_id: Logprob(logprob=-0.1),
50 }
51 ],
52 finish_reason=None,
53 )
54 ],
55 finished=True,
56 )
57
58
59 class _Serving(BeamSearchOnlineMixin):
60 renderer = _Renderer()
61 engine_client = _EngineClient()
62
63
64 @pytest.mark.asyncio
65 async def test_beam_search_handles_extra_logprob_candidates() -> None:
66 prompt: TokensInput = {
67 "type": "token",
68 "prompt": "prompt",
69 "prompt_token_ids": [1],
70 }
71 params = BeamSearchParams(beam_width=2, max_tokens=1)
72
73 outputs = [
74 output async for output in _Serving().beam_search(prompt, "request", params)
75 ]
76
77 assert len(outputs) == 1
78 assert outputs[0].outputs[0].finish_reason == "stop"
79 assert outputs[0].outputs[0].token_ids == []
80 assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1)
81
82
83 @pytest.mark.asyncio
84 @pytest.mark.parametrize(
85 ("skip_special_tokens", "expected_text"),
86 [
87 pytest.param(True, "", id="skip"),
88 pytest.param(False, "0", id="keep"),
89 ],
90 )
91 async def test_beam_search_respects_skip_special_tokens(
92 skip_special_tokens: bool, expected_text: str
93 ) -> None:
94 prompt: TokensInput = {
95 "type": "token",
96 "prompt": "prompt",
97 "prompt_token_ids": [1],
98 }
99 params = BeamSearchParams(
100 beam_width=1,
101 max_tokens=1,
102 ignore_eos=True,
103 skip_special_tokens=skip_special_tokens,
104 )
105
106 outputs = [
107 output async for output in _Serving().beam_search(prompt, "request", params)
108 ]
109
110 assert outputs[0].outputs[0].text == expected_text
111 assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id]
112
113
114 @pytest.mark.asyncio
115 @pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)])
116 @pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]])
117 async def test_beam_search_abort_returns_partial_outputs(
118 monkeypatch,
119 abort_after: int,
120 prompt_token: int,
121 terminal_logprobs: SampleLogprobs | None,
122 ) -> None:
123 """Abort returns each beam's tokens and scores from the last completed step."""
124 calls = 0
125
126 async def generate(prompt, *args, **kwargs):
127 nonlocal calls
128 calls += 1
129 depth = len(prompt["prompt_token_ids"]) - 1
130 assert depth <= abort_after, "Beam search continued after abort"
131 result = await anext(_EngineClient().generate(prompt, *args, **kwargs))
132 output = result.outputs[0]
133 output.token_ids = [11]
134 output.logprobs = [{11: Logprob(-0.1), 12: Logprob(-0.2)}]
135 output.finish_reason = "length"
136 if depth == abort_after and prompt["prompt_token_ids"][-1] != 12:
137 output.token_ids = []
138 output.logprobs = terminal_logprobs
139 output.finish_reason = "abort"
140 yield result
141
142 serving = _Serving()
143 monkeypatch.setattr(serving.engine_client, "generate", generate)
144 prompt: TokensInput = {
145 "type": "token",
146 "prompt": "prompt",
147 "prompt_token_ids": [prompt_token],
148 }
149 outputs = [
150 output
151 async for output in serving.beam_search(
152 prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4)
153 )
154 ]
155
156 assert calls == 1 + 2 * abort_after
157 assert len(outputs) == 1
158 assert outputs[0].finished
159 assert outputs[0].request_id == "request"
160 expected_tokens = [[11], [12]] if abort_after else [[]]
161 assert [output.token_ids for output in outputs[0].outputs] == expected_tokens
162 for output in outputs[0].outputs:
163 assert output.finish_reason == "abort"
164 assert output.text == _Tokenizer().decode(list(output.token_ids))
165 assert output.logprobs is not None
166 assert len(output.logprobs) == abort_after
167 assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after)