@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 210 lines 7.93 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """Unit tests for GPUModelRunner._update_streaming_request function."""
5
6 from unittest.mock import Mock
7
8 import pytest
9
10 from vllm.multimodal.inputs import (
11 MultiModalFeatureSpec,
12 MultiModalKwargsItem,
13 PlaceholderRange,
14 )
15 from vllm.sampling_params import SamplingParams
16 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
17 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
18
19 # Not cpu_test: InputBatch allocates pinned (UVA) memory, which requires a
20 # CUDA device even though the batch tensors live on the CPU.
21
22
23 @pytest.fixture
24 def mock_model_runner_with_input_batch():
25 """Create a mock GPUModelRunner with a real InputBatch for e2e testing."""
26 runner = Mock(spec=GPUModelRunner)
27 runner.uses_mrope = False
28 runner.requests = {}
29 runner.max_num_reqs = 10
30 runner.max_model_len = 1024
31 runner.late_interaction_runner = Mock()
32
33 # Create a real InputBatch for e2e testing
34 runner.input_batch = InputBatch(
35 max_num_reqs=10,
36 max_model_len=1024,
37 max_num_batched_tokens=1024,
38 device="cpu",
39 vocab_size=32000,
40 block_sizes=[16],
41 kernel_block_sizes=[16],
42 max_num_blocks_per_req=[64],
43 logitsprocs=None,
44 is_pooling_model=False,
45 )
46 return runner
47
48
49 def test_e2e_streaming_request_update_basic_flow(mock_model_runner_with_input_batch):
50 """Test that streaming session are updated correctly.
51
52 This test validates that when a streaming session is updated with new prompt tokens:
53 1. The request is removed from InputBatch before updating (avoids duplication)
54 2. Request state fields are updated correctly
55 3. output_token_ids is cleared (intermediate outputs are now in prompt_token_ids)
56 """
57 runner = mock_model_runner_with_input_batch
58 req_id = "streaming_req_0"
59
60 # Step 1: Create initial request state with some computed tokens
61 initial_req_state = CachedRequestState(
62 req_id=req_id,
63 prompt_token_ids=[1, 2, 3],
64 mm_features=[],
65 sampling_params=SamplingParams(temperature=0.5),
66 pooling_params=None,
67 generator=None,
68 block_ids=([0],),
69 num_computed_tokens=3,
70 output_token_ids=[10, 11], # Generated 2 tokens
71 )
72 runner.requests[req_id] = initial_req_state
73
74 # Add request to InputBatch
75 runner.input_batch.add_request(initial_req_state)
76 assert req_id in runner.input_batch.req_id_to_index
77
78 # Step 2: Create new request data with extended prompt
79 # The scheduler has already set prompt_token_ids to the full sequence
80 # (original prompt + intermediate outputs + new prompt)
81 new_req_data = Mock()
82 new_req_data.prompt_token_ids = [
83 1,
84 2,
85 3,
86 10,
87 4,
88 5,
89 ] # Full sequence with intermediate output (10)
90 new_req_data.mm_features = []
91 new_req_data.prompt_embeds = None
92 new_req_data.sampling_params = SamplingParams(temperature=0.8, max_tokens=50)
93 new_req_data.pooling_params = None
94 new_req_data.block_ids = ([0, 1],)
95 new_req_data.num_computed_tokens = 4 # 3 original prompt + 1 intermediate output
96
97 # Step 3: Update the request
98 updated_req_state = GPUModelRunner._update_streaming_request(
99 runner, req_id, new_req_data
100 )
101
102 # Step 4: Verify the request state was updated correctly
103 assert updated_req_state.prompt_token_ids == [1, 2, 3, 10, 4, 5]
104 assert updated_req_state.num_computed_tokens == 4
105 assert updated_req_state.sampling_params.temperature == 0.8
106 assert updated_req_state.sampling_params.max_tokens == 50
107 assert updated_req_state.block_ids == ([0, 1],)
108
109 # Verify output_token_ids were cleared
110 # (intermediate outputs are now in prompt_token_ids)
111 assert updated_req_state.output_token_ids == []
112
113 # Verify the same object is returned
114 assert runner.requests[req_id] is updated_req_state
115
116 # Verify request was removed from InputBatch during update (avoids duplication)
117 assert req_id not in runner.input_batch.req_id_to_index
118
119
120 def test_e2e_streaming_with_multimodal_features(mock_model_runner_with_input_batch):
121 """Test that streaming session with multimodal features are updated correctly.
122
123 This test validates that when a streaming session with mm features is updated:
124 1. The request is removed from InputBatch before updating (avoids duplication)
125 2. Multimodal features from both requests are preserved and merged correctly
126 3. New prompt tokens (including intermediate outputs) are appended correctly
127 4. output_token_ids is cleared (intermediate outputs are now in prompt_token_ids)
128 """
129 runner = mock_model_runner_with_input_batch
130 req_id = "streaming_mm_req_0"
131
132 # Step 1: Create initial request state with one multimodal feature
133 mm_feature_1 = MultiModalFeatureSpec(
134 data=MultiModalKwargsItem.dummy(),
135 modality="audio",
136 identifier="audio_1",
137 mm_position=PlaceholderRange(offset=2, length=10),
138 )
139
140 initial_req_state = CachedRequestState(
141 req_id=req_id,
142 prompt_token_ids=[1, 2] + [0] * 10 + [3, 4], # 2 + 10 (mm) + 2 = 14 tokens
143 mm_features=[mm_feature_1],
144 sampling_params=SamplingParams(),
145 pooling_params=None,
146 generator=None,
147 block_ids=([0],),
148 num_computed_tokens=14,
149 output_token_ids=[100], # Generated 1 token
150 )
151 runner.requests[req_id] = initial_req_state
152
153 # Add request to InputBatch
154 runner.input_batch.add_request(initial_req_state)
155 assert req_id in runner.input_batch.req_id_to_index
156
157 # Step 2: Create new request data with additional multimodal feature
158 # The scheduler has already set prompt_token_ids to the full sequence
159 # (original prompt + intermediate outputs + new prompt with new multimodal feature)
160 mm_feature_2 = MultiModalFeatureSpec(
161 data=MultiModalKwargsItem.dummy(),
162 modality="audio",
163 identifier="audio_2",
164 mm_position=PlaceholderRange(offset=15, length=5),
165 )
166
167 new_req_data = Mock()
168 # Full sequence: [1, 2] + [0]*10 + [3, 4] + [100] + [0]*5 + [5] = 21 tokens
169 new_req_data.prompt_token_ids = [1, 2] + [0] * 10 + [3, 4, 100] + [0] * 5 + [5]
170 new_req_data.mm_features = [mm_feature_1, mm_feature_2]
171 new_req_data.prompt_embeds = None
172 new_req_data.sampling_params = SamplingParams(temperature=0.7, max_tokens=30)
173 new_req_data.pooling_params = None
174 new_req_data.block_ids = ([0, 1],)
175 new_req_data.num_computed_tokens = 14 # 14 tokens from initial request
176
177 # Step 3: Update the request
178 updated_req_state = GPUModelRunner._update_streaming_request(
179 runner, req_id, new_req_data
180 )
181
182 # Step 4: Verify the request state was updated correctly
183 # Verify multimodal features are preserved
184 assert len(updated_req_state.mm_features) == 2
185 assert updated_req_state.mm_features[0] == mm_feature_1
186 assert updated_req_state.mm_features[1] == mm_feature_2
187
188 # Verify prompt tokens include intermediate output (100) and new tokens
189 # Initial: 2 + 10 (mm1) + 2 = 14 tokens
190 # New: 2 + 10 (mm1) + 2 + 1 (output 100) + 5 (mm2) + 1 = 21 tokens
191 assert len(updated_req_state.prompt_token_ids) == 21
192 assert updated_req_state.prompt_token_ids == [1, 2] + [0] * 10 + [3, 4, 100] + [
193 0
194 ] * 5 + [5]
195
196 # Verify output_token_ids were cleared
197 # (intermediate outputs are now in prompt_token_ids)
198 assert updated_req_state.output_token_ids == []
199
200 # Verify other parameters were updated
201 assert updated_req_state.num_computed_tokens == 14
202 assert updated_req_state.sampling_params.temperature == 0.7
203 assert updated_req_state.sampling_params.max_tokens == 30
204 assert updated_req_state.block_ids == ([0, 1],)
205
206 # Verify the same object is returned
207 assert runner.requests[req_id] is updated_req_state
208
209 # Verify request was removed from InputBatch during update (avoids duplication)
210 assert req_id not in runner.input_batch.req_id_to_index