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, 07:00 UTC
[Model Runner V2] Support Streaming Inputs (#37028)
Signed-off-by: Santino Ramos <elsantinoramos@gmail.com>
Santino Ramos committed
Mar 20, 2026 at 13:42 UTC
85f671b8e1cfa6655605efdf1263cf2f94e9b992
7 files changed
+263
-12
tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py
new
+207
@@ -0,0 +1,207 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Unit tests for MRv2 GPUModelRunner.add_requests streaming input support."""
5
+
6
+from unittest.mock import Mock
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.v1.core.sched.output import (
12
+ CachedRequestData,
13
+ NewRequestData,
14
+ SchedulerOutput,
15
+)
16
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
17
+from vllm.v1.worker.gpu.states import RequestState
18
+
19
+pytestmark = pytest.mark.cpu_test
20
+
21
+
22
+@pytest.fixture
23
+def mock_model_runner_with_req_states():
24
+ """Create a mock MRv2 GPUModelRunner with a real RequestState."""
25
+
26
+ runner = Mock(spec=GPUModelRunner)
27
+ runner.req_states = RequestState(
28
+ max_num_reqs=10,
29
+ max_model_len=1024,
30
+ max_num_batched_tokens=1024,
31
+ num_speculative_steps=0,
32
+ vocab_size=32000,
33
+ device=torch.device("cpu"),
34
+ model_dtype=torch.float32,
35
+ cache_draft_logits=False,
36
+ )
37
+ runner.encoder_cache = None
38
+ runner.model_state = Mock()
39
+ runner.block_tables = Mock()
40
+ runner.lora_state = Mock()
41
+ runner.sampler = None
42
+ runner.prompt_logprobs_worker = None
43
+ runner.is_last_pp_rank = False
44
+
45
+ # Mock staged writes — they use Triton kernels that require GPU
46
+ runner.req_states.apply_staged_writes = Mock()
47
+
48
+ # Bind the real methods to our mock
49
+ runner._remove_request = GPUModelRunner._remove_request.__get__(runner)
50
+ runner.add_requests = GPUModelRunner.add_requests.__get__(runner)
51
+ return runner
52
+
53
+
54
+def _make_scheduler_output(new_reqs):
55
+ return SchedulerOutput(
56
+ scheduled_new_reqs=new_reqs,
57
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
58
+ num_scheduled_tokens={},
59
+ total_num_scheduled_tokens=0,
60
+ scheduled_spec_decode_tokens={},
61
+ scheduled_encoder_inputs={},
62
+ num_common_prefix_blocks=[],
63
+ finished_req_ids=set(),
64
+ free_encoder_mm_hashes=[],
65
+ )
66
+
67
+
68
+def test_e2e_streaming_request_update_basic_flow(
69
+ mock_model_runner_with_req_states,
70
+):
71
+ """Test that streaming sessions are updated correctly.
72
+
73
+ This test validates that when a streaming session is updated with new
74
+ prompt tokens:
75
+ 1. The old request state is removed (no free_indices leak)
76
+ 2. The new state is written with updated prefill_token_ids
77
+ 3. model_state and block_tables are re-registered for the new state
78
+ """
79
+ runner = mock_model_runner_with_req_states
80
+ req_states = runner.req_states
81
+ req_id = "streaming_req_0"
82
+ initial_free = len(req_states.free_indices)
83
+
84
+ # Step 1: Add initial request with 3 prompt tokens, all computed
85
+ initial_req_data = NewRequestData(
86
+ req_id=req_id,
87
+ prompt_token_ids=[1, 2, 3],
88
+ prefill_token_ids=[1, 2, 3],
89
+ mm_features=[],
90
+ sampling_params=None,
91
+ pooling_params=None,
92
+ block_ids=([0],),
93
+ num_computed_tokens=3,
94
+ lora_request=None,
95
+ )
96
+ runner.add_requests(_make_scheduler_output([initial_req_data]))
97
+ assert req_id in req_states.req_id_to_index
98
+ assert len(req_states.free_indices) == initial_free - 1
99
+
100
+ # Step 2: Create streaming update with extended prompt
101
+ # The scheduler has already set prefill_token_ids to the full sequence
102
+ # (original prompt + intermediate output + new prompt tokens)
103
+ updated_req_data = NewRequestData(
104
+ req_id=req_id,
105
+ prompt_token_ids=[1, 2, 3],
106
+ prefill_token_ids=[1, 2, 3, 10, 4, 5],
107
+ mm_features=[],
108
+ sampling_params=None,
109
+ pooling_params=None,
110
+ block_ids=([0, 1],),
111
+ num_computed_tokens=4, # 3 original prompt + 1 intermediate output
112
+ lora_request=None,
113
+ )
114
+ runner.add_requests(_make_scheduler_output([updated_req_data]))
115
+
116
+ # Step 3: Verify no free_indices leak (old slot recycled)
117
+ assert len(req_states.free_indices) == initial_free - 1
118
+
119
+ # Verify the request is still tracked with exactly one index
120
+ assert req_id in req_states.req_id_to_index
121
+ assert sum(1 for v in req_states.index_to_req_id.values() if v == req_id) == 1
122
+
123
+ # Verify state was updated with new values
124
+ new_idx = req_states.req_id_to_index[req_id]
125
+ assert req_states.prompt_len.np[new_idx] == 3
126
+ assert req_states.prefill_len.np[new_idx] == 6
127
+ assert req_states.num_computed_prefill_tokens[new_idx] == 4
128
+
129
+ # Verify model_state and block_tables were re-registered
130
+ runner.model_state.add_request.assert_called_with(new_idx, updated_req_data)
131
+ runner.block_tables.append_block_ids.assert_called_with(
132
+ new_idx, ([0, 1],), overwrite=True
133
+ )
134
+
135
+
136
+def test_e2e_streaming_with_multimodal_features(
137
+ mock_model_runner_with_req_states,
138
+):
139
+ """Test that streaming sessions with multimodal features are updated.
140
+
141
+ This test validates that when a streaming session with mm features
142
+ is updated:
143
+ 1. The old request state is removed (no free_indices leak)
144
+ 2. encoder_cache is cleaned up and re-registered with new mm_features
145
+ 3. model_state is re-registered (recomputes M-RoPE positions etc.)
146
+ """
147
+ runner = mock_model_runner_with_req_states
148
+ req_states = runner.req_states
149
+ req_id = "streaming_mm_req_0"
150
+ initial_free = len(req_states.free_indices)
151
+
152
+ # Enable encoder_cache for multimodal
153
+ runner.encoder_cache = Mock()
154
+
155
+ # Step 1: Add initial request with one audio feature
156
+ mm_feature_1 = Mock()
157
+ initial_req_data = NewRequestData(
158
+ req_id=req_id,
159
+ prompt_token_ids=[1, 2] + [0] * 10 + [3, 4],
160
+ prefill_token_ids=[1, 2] + [0] * 10 + [3, 4],
161
+ mm_features=[mm_feature_1],
162
+ sampling_params=None,
163
+ pooling_params=None,
164
+ block_ids=([0],),
165
+ num_computed_tokens=14,
166
+ lora_request=None,
167
+ )
168
+ runner.add_requests(_make_scheduler_output([initial_req_data]))
169
+ assert req_id in req_states.req_id_to_index
170
+
171
+ # Reset mocks to track only the streaming update calls
172
+ runner.encoder_cache.reset_mock()
173
+ runner.model_state.reset_mock()
174
+
175
+ # Step 2: Create streaming update with additional multimodal feature
176
+ # The scheduler has folded the intermediate output (100) into
177
+ # prefill_token_ids and added a new audio chunk
178
+ mm_feature_2 = Mock()
179
+ updated_req_data = NewRequestData(
180
+ req_id=req_id,
181
+ prompt_token_ids=[1, 2] + [0] * 10 + [3, 4],
182
+ prefill_token_ids=[1, 2] + [0] * 10 + [3, 4, 100] + [0] * 5 + [5],
183
+ mm_features=[mm_feature_1, mm_feature_2],
184
+ sampling_params=None,
185
+ pooling_params=None,
186
+ block_ids=([0, 1],),
187
+ num_computed_tokens=14,
188
+ lora_request=None,
189
+ )
190
+ runner.add_requests(_make_scheduler_output([updated_req_data]))
191
+
192
+ # Step 3: Verify no free_indices leak
193
+ assert len(req_states.free_indices) == initial_free - 1
194
+ assert sum(1 for v in req_states.index_to_req_id.values() if v == req_id) == 1
195
+
196
+ # Verify encoder_cache was cleaned up and re-registered
197
+ runner.encoder_cache.remove_request.assert_called_once_with(req_id)
198
+ runner.encoder_cache.add_request.assert_called_once_with(
199
+ req_id, [mm_feature_1, mm_feature_2]
200
+ )
201
+
202
+ # Verify model_state was re-registered with new data
203
+ new_idx = req_states.req_id_to_index[req_id]
204
+ runner.model_state.add_request.assert_called_once_with(new_idx, updated_req_data)
205
+
206
+ # Verify updated prefill length
207
+ assert req_states.prefill_len.np[new_idx] == 21
vllm/model_executor/models/whisper_causal.py
+4
-2
@@ -150,8 +150,10 @@ def create_whisper_attention_backend_with_block_pooling(
150
new_common_attn_metadata.query_start_loc *= block_pool_size
151
new_common_attn_metadata.query_start_loc_cpu *= block_pool_size
152
new_common_attn_metadata.seq_lens *= block_pool_size
153
- new_common_attn_metadata._seq_lens_cpu *= block_pool_size
154
- new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size
153
+ if new_common_attn_metadata._seq_lens_cpu is not None:
154
+ new_common_attn_metadata._seq_lens_cpu *= block_pool_size
155
+ if new_common_attn_metadata._num_computed_tokens_cpu is not None:
156
+ new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size
157
new_common_attn_metadata.num_actual_tokens *= block_pool_size
158
new_common_attn_metadata.max_query_len *= block_pool_size
159
new_common_attn_metadata.max_seq_len *= block_pool_size
vllm/v1/worker/gpu/attn_utils.py
+6
-1
@@ -111,6 +111,7 @@ def _reshape_kv_cache(
111
kv_cache_config: KVCacheConfig,
112
kv_cache_raw_tensors: dict[str, torch.Tensor],
113
attn_backends: dict[str, AttentionBackend],
114
+ cache_dtype: str,
115
) -> dict[str, torch.Tensor]:
116
kv_caches: dict[str, torch.Tensor] = {}
117
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
@@ -127,6 +128,7 @@ def _reshape_kv_cache(
128
kv_cache_spec.block_size,
129
kv_cache_spec.num_kv_heads,
130
kv_cache_spec.head_size,
131
+ cache_dtype,
132
)
133
134
# FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
@@ -155,9 +157,12 @@ def init_kv_cache(
157
kv_cache_config: KVCacheConfig,
158
attn_backends: dict[str, AttentionBackend],
159
device: torch.device,
160
+ cache_dtype: str,
161
) -> dict[str, torch.Tensor]:
162
kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
160
- kv_caches = _reshape_kv_cache(kv_cache_config, kv_cache_raw_tensors, attn_backends)
163
+ kv_caches = _reshape_kv_cache(
164
+ kv_cache_config, kv_cache_raw_tensors, attn_backends, cache_dtype
165
+ )
166
bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
167
return kv_caches
168
vllm/v1/worker/gpu/model_runner.py
+18
-6
@@ -359,6 +359,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
359
self.kv_cache_config,
360
self.attn_backends,
361
self.device,
362
+ self.cache_config.cache_dtype,
363
)
364
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
365
@@ -555,18 +556,23 @@ class GPUModelRunner(LoRAModelRunnerMixin):
556
)
557
return cuda_graph_size
558
559
+ def _remove_request(self, req_id: str) -> bool:
560
+ if not self.req_states.remove_request(req_id):
561
+ return False
562
+ if self.encoder_cache is not None:
563
+ self.encoder_cache.remove_request(req_id)
564
+ if self.prompt_logprobs_worker is not None:
565
+ self.prompt_logprobs_worker.remove_request(req_id)
566
+ self.lora_state.remove_request(req_id)
567
+ return True
568
+
569
def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
570
finished_req_ids = scheduler_output.finished_req_ids
571
preempted_req_ids = scheduler_output.preempted_req_ids
572
if preempted_req_ids:
573
finished_req_ids = finished_req_ids.union(preempted_req_ids)
574
for req_id in finished_req_ids:
564
- self.req_states.remove_request(req_id)
565
- if self.encoder_cache is not None:
566
- self.encoder_cache.remove_request(req_id)
567
- if self.prompt_logprobs_worker is not None:
568
- self.prompt_logprobs_worker.remove_request(req_id)
569
- self.lora_state.remove_request(req_id)
575
+ self._remove_request(req_id)
576
577
def free_states(self, scheduler_output: SchedulerOutput) -> None:
578
if self.encoder_cache is not None:
@@ -578,6 +584,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
584
assert new_req_data.prompt_token_ids is not None
585
assert new_req_data.prefill_token_ids is not None
586
req_id = new_req_data.req_id
587
+
588
+ # Streaming input update: request already exists from a prior
589
+ # chunk. Remove old state so it can be cleanly re-added below
590
+ # with the updated prompt_token_ids and mm_features.
591
+ self._remove_request(req_id)
592
+
593
prompt_len = len(new_req_data.prompt_token_ids)
594
self.req_states.add_request(
595
req_id=req_id,
vllm/v1/worker/gpu/model_states/default.py
+23
@@ -7,6 +7,7 @@ import torch.nn as nn
7
8
from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
+from vllm.tasks import GenerationTask
11
from vllm.v1.core.sched.output import NewRequestData
12
from vllm.v1.kv_cache_interface import KVCacheConfig
13
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
@@ -61,6 +62,28 @@ class DefaultModelState(ModelState):
62
device=self.device,
63
)
64
65
+ def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
66
+ from vllm.model_executor.models.interfaces import (
67
+ supports_realtime,
68
+ supports_transcription,
69
+ )
70
+ from vllm.model_executor.models.interfaces_base import is_text_generation_model
71
+
72
+ supported_tasks = list[GenerationTask]()
73
+
74
+ if is_text_generation_model(self.model):
75
+ supported_tasks.append("generate")
76
+
77
+ if supports_transcription(self.model):
78
+ if self.model.supports_transcription_only:
79
+ return ("transcription",)
80
+ supported_tasks.append("transcription")
81
+
82
+ if supports_realtime(self.model):
83
+ supported_tasks.append("realtime")
84
+
85
+ return tuple(supported_tasks)
86
+
87
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
88
if self.rope_state is not None:
89
assert new_req_data.prefill_token_ids is not None
vllm/v1/worker/gpu/model_states/interface.py
+2
-1
@@ -28,8 +28,9 @@ class ModelState(ABC):
28
) -> None:
29
raise NotImplementedError
30
31
+ @abstractmethod
32
def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
32
- return ("generate",)
33
+ raise NotImplementedError
34
35
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
36
return None
vllm/v1/worker/gpu/states.py
+3
-2
@@ -109,13 +109,14 @@ class RequestState:
109
self.all_token_ids.apply_write()
110
self.num_computed_tokens.apply_write()
111
112
- def remove_request(self, req_id: str) -> None:
112
+ def remove_request(self, req_id: str) -> bool:
113
req_idx = self.req_id_to_index.pop(req_id, None)
114
if req_idx is None:
115
# Request not found.
116
- return
116
+ return False
117
self.index_to_req_id.pop(req_idx, None)
118
self.free_indices.append(req_idx)
119
+ return True
120
121
def any_prefills(self, idx_mapping_np: np.ndarray) -> bool:
122
return np.any(