@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 260 lines 9.08 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import copy
4
5 import pytest
6
7 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
8 from vllm.v1.request import FinishReason, RequestStatus
9
10 from .utils import (
11 assert_scheduler_empty,
12 create_model_runner_output,
13 create_request,
14 create_scheduler,
15 create_vllm_config,
16 )
17
18 pytestmark = pytest.mark.cpu_test
19
20
21 def test_basic_lifecycle():
22 """Test lifecycle of a Remote Decode request."""
23 vllm_config = create_vllm_config()
24 scheduler = create_scheduler(vllm_config)
25
26 # 2 Full Blocks and 1 Half Block.
27 BLOCK_SIZE = vllm_config.cache_config.block_size
28 NUM_EXTERNAL_FULL_BLOCKS = 2
29 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
30
31 request = create_request(
32 request_id=1,
33 block_size=BLOCK_SIZE,
34 max_tokens=1,
35 num_tokens=NUM_TOKENS,
36 do_remote_decode=True,
37 )
38
39 scheduler.add_request(request)
40 request_id = request.request_id
41
42 # STEP (1): Prefill.
43 # (1a): schedule()
44 scheduler_output = scheduler.schedule()
45 assert len(scheduler.requests) == 1
46 assert len(scheduler.running) == 1
47 assert len(scheduler_output.scheduled_new_reqs) == 1
48
49 # (1b): execute_model()
50 model_runner_output = create_model_runner_output(reqs=[request])
51
52 # (1c): update_from_output()
53 engine_core_outputs = scheduler.update_from_output(
54 scheduler_output, model_runner_output
55 )
56
57 # Ensure the request is finished after 1 token.
58 assert request.is_finished()
59 assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
60 output = engine_core_outputs[0].outputs[0]
61 assert output.finish_reason == FinishReason.LENGTH
62 assert output.kv_transfer_params is not None
63
64 # Request freed in Scheduler and in Persistent Batch ...
65 assert request_id in scheduler.finished_req_ids
66 assert len(scheduler.running) == 0
67 assert len(scheduler.waiting) == 0
68
69 # ... but blocks should not be freed.
70 assert len(scheduler.requests) == 1
71 blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
72 0
73 ].req_to_blocks[request_id]
74 for block in blocks:
75 assert block.ref_cnt == 1
76
77 # STEP (2): Send Finished to PB.
78 # (2a): schedule() - pass finished request to PB.
79 scheduler_output = scheduler.schedule()
80 assert len(scheduler.requests) == 1
81 assert len(scheduler.running) == 0
82 assert len(scheduler_output.finished_req_ids) == 1
83 assert request_id in scheduler_output.finished_req_ids
84 assert len(scheduler_output.scheduled_new_reqs) == 0
85 assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
86 assert len(scheduler.finished_req_ids) == 0
87
88 # (2b): execute_model()
89 model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
90
91 # (2c): update_from_output()
92 scheduler.update_from_output(scheduler_output, model_runner_output)
93
94 # STEP (3): Finished sending.
95 # (3a): schedule() - pass finished request to PB.
96 scheduler_output = scheduler.schedule()
97 assert len(scheduler.requests) == 1
98 assert len(scheduler.running) == 0
99 assert len(scheduler_output.finished_req_ids) == 0
100 assert len(scheduler_output.scheduled_new_reqs) == 0
101 assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
102 assert len(scheduler.finished_req_ids) == 0
103
104 # (3b): execute_model()
105 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
106 model_runner_output.kv_connector_output = KVConnectorOutput(
107 finished_sending={request_id}
108 )
109
110 # (3c): update_from_output()
111 scheduler.update_from_output(scheduler_output, model_runner_output)
112
113 # Confirm we do not have any memory leaks after req lifecycle.
114 assert_scheduler_empty(scheduler)
115
116
117 def test_short_prompt_lifecycle():
118 """Test lifecycle of a Remote Decode request with short prompt."""
119 vllm_config = create_vllm_config()
120 scheduler = create_scheduler(vllm_config)
121
122 # Not enough tokens for full block.
123 BLOCK_SIZE = vllm_config.cache_config.block_size
124 NUM_TOKENS = BLOCK_SIZE // 2
125 request = create_request(
126 request_id=1,
127 block_size=BLOCK_SIZE,
128 max_tokens=1,
129 num_tokens=NUM_TOKENS,
130 do_remote_decode=True,
131 )
132
133 scheduler.add_request(request)
134
135 # STEP (1): Prefill.
136 # (1a): schedule()
137 scheduler_output = scheduler.schedule()
138 assert len(scheduler.requests) == 1
139 assert len(scheduler.running) == 1
140 assert len(scheduler_output.scheduled_new_reqs) == 1
141
142 # (1b): execute_model()
143 model_runner_output = create_model_runner_output(reqs=[request])
144
145 # (1c): update_from_output()
146 # Even though tokens < block_size, there will be kv xfer for partial block.
147 eco = scheduler.update_from_output(scheduler_output, model_runner_output)
148 kv_transfer_params = eco[0].outputs[0].kv_transfer_params
149
150 assert len(kv_transfer_params["remote_block_ids"]) == 1
151
152 # Confirm we do not have any memory leaks after req lifecycle.
153 # We need to mark sending finish to clear data for persistent batch.
154 scheduler_output = scheduler.schedule()
155 # Use create_model_runner_output to pass kv_connector_output along
156 model_runner_output = create_model_runner_output(
157 reqs=[request], finished_sending={request.request_id}
158 )
159 scheduler.update_from_output(scheduler_output, model_runner_output)
160 assert_scheduler_empty(scheduler)
161
162
163 def test_prefix_cache_lifecycle():
164 """Test that remote decode params still work with a prefix cache hit."""
165 vllm_config = create_vllm_config()
166 scheduler = create_scheduler(vllm_config)
167
168 # Prime the KVCache.
169 BLOCK_SIZE = vllm_config.cache_config.block_size
170 NUM_EXTERNAL_FULL_BLOCKS = 3
171 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
172
173 request_normal = create_request(
174 request_id=1, block_size=BLOCK_SIZE, num_tokens=NUM_TOKENS
175 )
176
177 scheduler.add_request(request_normal)
178 scheduler_output = scheduler.schedule()
179 model_runner_output = create_model_runner_output(
180 reqs=[request_normal], use_eos=True
181 )
182 scheduler.update_from_output(scheduler_output, model_runner_output)
183 scheduler_output = scheduler.schedule()
184 scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
185
186 #####################
187 # Actual Test: confirm we send all blocks.
188
189 # Step (1): Send the KV Transfer.
190 NUM_EXTERNAL_FULL_BLOCKS -= 1
191 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
192
193 request_remote = create_request(
194 request_id=1,
195 block_size=BLOCK_SIZE,
196 num_tokens=NUM_TOKENS,
197 do_remote_decode=True,
198 )
199
200 scheduler.add_request(request_remote)
201 scheduler_output = scheduler.schedule()
202 model_runner_output = create_model_runner_output(reqs=[request_remote])
203 eco = scheduler.update_from_output(scheduler_output, model_runner_output)
204 kv_transfer_params = eco[0].outputs[0].kv_transfer_params
205
206 # Ensure we send all block ids, including the partial blocks,
207 # even if there is a cache hit.
208 # remote_block_ids is BlockIds (tuple of lists); sum block counts across groups.
209 num_remote_blocks = sum(len(g) for g in kv_transfer_params["remote_block_ids"])
210 assert num_remote_blocks == (NUM_EXTERNAL_FULL_BLOCKS + 1)
211
212 # STEP (2): Ensure it is freed.
213 scheduler_output = scheduler.schedule()
214 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
215 model_runner_output.kv_connector_output = KVConnectorOutput(
216 finished_sending={request_remote.request_id}
217 )
218 scheduler.update_from_output(scheduler_output, model_runner_output)
219 assert_scheduler_empty(scheduler)
220
221
222 def test_abort_during_kv_transfer():
223 """Test aborting request does not release blocks for remote decode."""
224 vllm_config = create_vllm_config()
225 scheduler = create_scheduler(vllm_config)
226
227 # Prime the KVCache.
228 BLOCK_SIZE = vllm_config.cache_config.block_size
229 NUM_EXTERNAL_FULL_BLOCKS = 2
230 NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
231
232 request = create_request(
233 request_id=1,
234 block_size=BLOCK_SIZE,
235 num_tokens=NUM_TOKENS,
236 do_remote_decode=True,
237 )
238
239 scheduler.add_request(request)
240 scheduler_output = scheduler.schedule()
241 model_runner_output = create_model_runner_output(reqs=[request])
242 scheduler.update_from_output(scheduler_output, model_runner_output)
243 scheduler_output = scheduler.schedule()
244 scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
245
246 # Request removed from PB but blocks should not be freed.
247 assert len(scheduler.requests) == 1
248
249 # Abort the request, and check the blocks are still not freed
250 scheduler.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED)
251 assert len(scheduler.requests) == 1
252
253 # Simulate a finished sending notification
254 scheduler_output = scheduler.schedule()
255 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
256 model_runner_output.kv_connector_output = KVConnectorOutput(
257 finished_sending=[request.request_id]
258 )
259 scheduler.update_from_output(scheduler_output, model_runner_output)
260 assert_scheduler_empty(scheduler)