@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b7433ca1a4
Fork 0

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 14, 05:30 UTC

[Spec Decode] Efficient padded speculation (#24539)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>

Benjamin Chislett committed Sep 18, 2025 at 01:07 UTC b7433ca1a47732394b1bdea4099d98389515954b
5 files changed +507 -104
tests/v1/spec_decode/test_eagle.py
+174 -5
@@ -19,6 +19,8 @@ from vllm.config.load import LoadConfig
19 from vllm.model_executor.models.llama import LlamaForCausalLM
20 from vllm.platforms import current_platform
21 from vllm.v1.spec_decode.eagle import EagleProposer
22 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
23 +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
24
25 model_dir = "meta-llama/Llama-3.1-8B-Instruct"
26 eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
@@ -64,6 +66,86 @@ def _create_proposer(
66 device=current_platform.device_type)
67
68
69 +def test_prepare_next_token_ids():
70 + """
71 + Test for prepare_next_token_ids_cpu and prepare_next_token_ids_padded.
72 + Each will produce a device tensor of next_token_ids, taking as input
73 + either the GPU tensor of sampled_token_ids with -1 for rejected tokens,
74 + or the CPU python list[list[int]] with the rejected tokens removed.
75 + """
76 + device = torch.device(current_platform.device_type)
77 +
78 + num_requests = 4
79 + num_speculative_tokens = 4
80 + batch_spec = BatchSpec(
81 + seq_lens=[num_speculative_tokens + 1] * num_requests,
82 + query_lens=[num_speculative_tokens + 1] * num_requests,
83 + )
84 +
85 + req_ids = [f"req_{i+1}" for i in range(num_requests)]
86 + mock_input_batch = mock.MagicMock(spec=InputBatch)
87 + mock_input_batch.req_ids = req_ids
88 + mock_input_batch.num_reqs = num_requests
89 + mock_input_batch.vocab_size = 100
90 +
91 + mock_num_scheduled_tokens = {req_id: 0 for req_id in req_ids}
92 + mock_requests = {}
93 + for req_id in req_ids:
94 + mock_request = mock.MagicMock(spec=CachedRequestState)
95 + # Each request will have a backup next token id of 10, 20, 30, 40
96 + mock_request.get_token_id.return_value = int(req_id.split("_")[1]) * 10
97 + mock_request.num_computed_tokens = 0
98 + mock_requests[req_id] = mock_request
99 +
100 + sampled_token_ids = [
101 + [0, 1, -1, -1, -1], # 1 accepted, 3 rejected, "1" sampled
102 + [0, 1, 2, 3, 4], # all accepted, "4" sampled
103 + [-1, -1, -1, -1, -1], # sampling skipped, use backup token "30"
104 + [-1, -1, -1, -1, -1] # this request will be discarded
105 + ]
106 + sampled_token_ids_tensor = torch.tensor(sampled_token_ids,
107 + dtype=torch.int32,
108 + device=device)
109 + sampled_token_ids_cpu = [[i for i in seq if i != -1]
110 + for seq in sampled_token_ids]
111 +
112 + expected_next_token_ids_cpu = [1, 4, 30, 40]
113 + expected_next_token_ids_tensor = torch.tensor(expected_next_token_ids_cpu,
114 + dtype=torch.int32,
115 + device=device)
116 +
117 + proposer = _create_proposer("eagle", num_speculative_tokens)
118 +
119 + next_token_ids_from_cpu = proposer.prepare_next_token_ids_cpu(
120 + sampled_token_ids_cpu, mock_requests, mock_input_batch,
121 + mock_num_scheduled_tokens)
122 +
123 + assert torch.equal(next_token_ids_from_cpu, expected_next_token_ids_tensor)
124 +
125 + common_attn_metadata = create_common_attn_metadata(
126 + batch_spec,
127 + block_size=16,
128 + device=device,
129 + )
130 +
131 + discarded_req_indices = torch.tensor([3], dtype=torch.int64, device=device)
132 + num_discarded_reqs = 1
133 +
134 + expected_valid_sampled_tokens_count = torch.tensor([2, 5, 0, 0],
135 + dtype=torch.int32,
136 + device=device)
137 +
138 + next_token_ids_from_padded, valid_sampled_tokens_count = \
139 + proposer.prepare_next_token_ids_padded(
140 + common_attn_metadata, sampled_token_ids_tensor, mock_requests,
141 + mock_input_batch, discarded_req_indices, num_discarded_reqs)
142 +
143 + assert torch.equal(next_token_ids_from_padded,
144 + expected_next_token_ids_tensor)
145 + assert torch.equal(valid_sampled_tokens_count,
146 + expected_valid_sampled_tokens_count)
147 +
148 +
149 def test_prepare_inputs():
150 """
151 cu_target_query_lens: [0, a, a + b, a + b + c]
@@ -90,10 +172,24 @@ def test_prepare_inputs():
172 device=device,
173 )
174
93 - # Rejected tokens per request: [1, 3, 2]
94 - num_rejected_tokens = torch.tensor([1, 3, 2],
95 - dtype=torch.int32,
96 - device=device)
175 + # If there are `k` sampled tokens, then `k-1` tokens are draft tokens
176 + # from the previous iteration, and the last token is the bonus token sampled
177 + # from the base model.
178 + num_draft_tokens = [3, 6, 4] # one less than query_lens
179 + # num rejected tokens is [1, 3, 2]
180 + ACCEPT_TOKEN = 0
181 + BONUS_TOKEN = 1
182 + REJECT_TOKEN = -1
183 + sampled_token_ids = [
184 + [ACCEPT_TOKEN, ACCEPT_TOKEN, REJECT_TOKEN, BONUS_TOKEN],
185 + [
186 + ACCEPT_TOKEN, ACCEPT_TOKEN, ACCEPT_TOKEN, REJECT_TOKEN,
187 + REJECT_TOKEN, REJECT_TOKEN, BONUS_TOKEN
188 + ],
189 + [ACCEPT_TOKEN, ACCEPT_TOKEN, REJECT_TOKEN, REJECT_TOKEN, BONUS_TOKEN]
190 + ]
191 + sampled_token_ids = [[i for i in seq if i != REJECT_TOKEN]
192 + for seq in sampled_token_ids]
193
194 # Expected calculations:
195 # query_len_per_req = [4, 7, 5]
@@ -125,7 +221,7 @@ def test_prepare_inputs():
221 proposer = _create_proposer("eagle", 1)
222
223 updated_metadata, token_indices = proposer.prepare_inputs(
128 - common_attn_metadata, num_rejected_tokens.cpu())
224 + common_attn_metadata, sampled_token_ids, num_draft_tokens)
225
226 assert torch.equal(updated_metadata.query_start_loc,
227 expected_cu_num_tokens)
@@ -133,6 +229,77 @@ def test_prepare_inputs():
229 assert torch.equal(token_indices, expected_token_indices)
230
231
232 +def test_prepare_inputs_padded():
233 + """
234 + Input scenario is 3 requests with num_speculative_tokens == 2 and:
235 + - Request 1: query_len = 3, rejected = 1
236 + - Request 2: query_len = 3, rejected = 0
237 + - Request 3: query_len = 3, rejected = 2
238 +
239 + Expected outputs:
240 + token_indices: [0, 1, 2,
241 + 3, 4, 5,
242 + 6, 7, 8]
243 + Reason: Deferred computation should not disturb the original indices.
244 +
245 + token_indices_to_sample: [1, 5, 6]
246 + Reason: After accounting for rejections, these are the valid token positions
247 + from the original indices to sample from.
248 + """
249 +
250 + device = torch.device(current_platform.device_type)
251 +
252 + expected_token_indices = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8],
253 + dtype=torch.int32,
254 + device=device)
255 + expected_token_indices_to_sample = torch.tensor([1, 5, 6],
256 + dtype=torch.int32,
257 + device=device)
258 +
259 + num_speculative_tokens = 2
260 + batch_spec = BatchSpec(
261 + seq_lens=[3, 3, 3],
262 + query_lens=[3, 3, 3],
263 + )
264 +
265 + common_attn_metadata = create_common_attn_metadata(
266 + batch_spec,
267 + block_size=16,
268 + device=device,
269 + )
270 +
271 + # Needed for cu_num_draft_tokens, which is expected to be [3, 6, 9]
272 + expected_query_start_loc = torch.tensor([0, 3, 6, 9],
273 + dtype=torch.int32,
274 + device=device)
275 + spec_decode_metadata = SpecDecodeMetadata.make_dummy(
276 + draft_token_ids=[[0] * num_speculative_tokens] * 3,
277 + device=device,
278 + )
279 +
280 + # num_rejected_tokens = [1, 0, 2]
281 + # num_draft_tokens = [2, 2, 2]
282 + # valid_sampled_tokens_count = num_draft_tokens + 1 - num_rejected_tokens
283 + valid_sampled_tokens_count = torch.tensor([2, 3, 1],
284 + dtype=torch.int32,
285 + device=device)
286 +
287 + proposer = _create_proposer("eagle", num_speculative_tokens)
288 +
289 + output_metadata, token_indices, token_indices_to_sample = \
290 + proposer.prepare_inputs_padded(
291 + common_attn_metadata,
292 + spec_decode_metadata,
293 + valid_sampled_tokens_count)
294 +
295 + assert output_metadata.max_query_len == 3
296 + assert torch.equal(output_metadata.query_start_loc,
297 + expected_query_start_loc)
298 + assert torch.equal(token_indices, expected_token_indices)
299 + assert torch.equal(token_indices_to_sample,
300 + expected_token_indices_to_sample)
301 +
302 +
303 @pytest.mark.parametrize("method", ["eagle", "eagle3"])
304 @pytest.mark.parametrize("attn_backend",
305 get_attn_backend_list_based_on_platform())
@@ -373,6 +540,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
540 target_positions=target_positions,
541 target_hidden_states=target_hidden_states,
542 next_token_ids=next_token_ids,
543 + last_token_indices=None,
544 common_attn_metadata=common_attn_metadata,
545 sampling_metadata=sampling_metadata)
546
@@ -526,6 +694,7 @@ def test_propose_tree(spec_token_tree):
694 target_positions=target_positions,
695 target_hidden_states=target_hidden_states,
696 next_token_ids=next_token_ids,
697 + last_token_indices=None,
698 common_attn_metadata=common_attn_metadata,
699 sampling_metadata=sampling_metadata)
700 assert result.shape == (batch_size, num_speculative_tokens)
vllm/config/speculative.py
+5
@@ -83,6 +83,11 @@ class SpeculativeConfig:
83 disable_by_batch_size: Optional[int] = None
84 """Disable speculative decoding for new incoming requests when the number
85 of enqueued requests is larger than this value, if provided."""
86 + disable_padded_drafter_batch: bool = False
87 + """Disable input padding for speculative decoding. If set to True,
88 + speculative input batches can contain sequences of different lengths,
89 + which may only be supported by certain attention backends. This currently
90 + only affects the EAGLE method of speculation."""
91
92 # Ngram proposer configuration
93 prompt_lookup_max: Optional[int] = None
vllm/v1/spec_decode/eagle.py
+223 -35
@@ -27,6 +27,9 @@ from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
27 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
28 from vllm.v1.kv_cache_interface import KVCacheConfig
29 from vllm.v1.sample.metadata import SamplingMetadata
30 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
31 +from vllm.v1.utils import CpuGpuBuffer
32 +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
33 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
34
35 logger = init_logger(__name__)
@@ -94,20 +97,26 @@ class EagleProposer:
97 dtype=self.dtype,
98 device=device)
99
100 + # We need +1 here because the arange is used to set query_start_loc,
101 + # which has one more element than batch_size.
102 max_batch_size = vllm_config.scheduler_config.max_num_seqs
98 - self.arange = torch.arange(
99 - # We need +1 here because the arange is used to set query_start_loc,
100 - # which has one more element than batch_size.
101 - max_batch_size + 1,
102 - device=device,
103 - dtype=torch.int32,
104 - )
103 + max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens)
104 + self.arange = torch.arange(max_num_slots_for_arange,
105 + device=device,
106 + dtype=torch.int32)
107
108 self.inputs_embeds = torch.zeros(
109 (self.max_num_tokens, self.hidden_size),
110 dtype=self.dtype,
111 device=device)
112
113 + self.backup_next_token_ids = CpuGpuBuffer(
114 + max_batch_size,
115 + dtype=torch.int32,
116 + pin_memory=is_pin_memory_available(),
117 + device=device,
118 + with_numpy=True)
119 +
120 # Determine allowed attention backends once during initialization.
121 self.allowed_attn_types: tuple[type[EagleAttentionMetadata], ...]
122 if current_platform.is_rocm():
@@ -156,13 +165,16 @@ class EagleProposer:
165 target_hidden_states: torch.Tensor,
166 # [batch_size]
167 next_token_ids: torch.Tensor,
168 + last_token_indices: Optional[torch.Tensor],
169 common_attn_metadata: CommonAttentionMetadata,
170 sampling_metadata: SamplingMetadata,
171 mm_embeds: Optional[list[torch.Tensor]] = None,
172 ) -> torch.Tensor:
173 num_tokens = target_token_ids.shape[0]
174 batch_size = next_token_ids.shape[0]
165 - last_token_indices = common_attn_metadata.query_start_loc[1:] - 1
175 +
176 + if last_token_indices is None:
177 + last_token_indices = common_attn_metadata.query_start_loc[1:] - 1
178
179 if self.method == "eagle3":
180 assert isinstance(self.model, Eagle3LlamaForCausalLM)
@@ -228,6 +240,12 @@ class EagleProposer:
240 last_hidden_states, hidden_states = ret_hidden_states
241 sample_hidden_states = last_hidden_states[last_token_indices]
242 logits = self.model.compute_logits(sample_hidden_states, None)
243 +
244 + # Early exit if there is only one draft token to be generated.
245 + if self.num_speculative_tokens == 1:
246 + draft_token_ids = logits.argmax(dim=-1)
247 + return draft_token_ids.view(-1, 1)
248 +
249 positions = target_positions[last_token_indices]
250 hidden_states = hidden_states[last_token_indices]
251
@@ -245,15 +263,12 @@ class EagleProposer:
263
264 draft_token_ids = logits.argmax(dim=-1)
265
248 - # Early exit if there is only one draft token to be generated.
249 - if self.num_speculative_tokens == 1:
250 - # [batch_size, 1]
251 - return draft_token_ids.view(-1, 1)
252 -
253 - # TODO: Currently, MTP module released by deepseek only has
254 - # one layer. Adapt this code to support multiple layers once
255 - # there's a multi-layer MTP module.
256 - assert isinstance(attn_metadata, self.allowed_attn_types)
266 + if not isinstance(attn_metadata, self.allowed_attn_types):
267 + raise ValueError(
268 + f"Unsupported attention metadata type for speculative "
269 + "decoding with num_speculative_tokens > 1: "
270 + f"{type(attn_metadata)}. Supported types are: "
271 + f"{self.allowed_attn_types}")
272
273 # Generate the remaining draft tokens.
274 draft_token_ids_list = [draft_token_ids]
@@ -263,10 +278,13 @@ class EagleProposer:
278 input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
279 else:
280 input_batch_size = batch_size
266 - attn_metadata.num_actual_tokens = batch_size
267 - attn_metadata.max_query_len = 1
268 - attn_metadata.query_start_loc = self.arange[:batch_size + 1]
269 - for _ in range(self.num_speculative_tokens - 1):
281 +
282 + common_attn_metadata.num_actual_tokens = batch_size
283 + common_attn_metadata.max_query_len = 1
284 + common_attn_metadata.query_start_loc = self.arange[:batch_size + 1]
285 + common_attn_metadata.query_start_loc_cpu = torch.from_numpy(
286 + self.token_arange_np[:batch_size + 1]).clone()
287 + for token_index in range(self.num_speculative_tokens - 1):
288 # Update the inputs.
289 # cast to int32 is crucial when eagle model is compiled.
290 # tensor.argmax() returns int64 by default.
@@ -286,27 +304,38 @@ class EagleProposer:
304 positions)
305
306 # Increment the sequence lengths.
289 - attn_metadata.max_seq_len += 1
290 - attn_metadata.seq_lens += 1
291 - # Consider max model length.
292 - attn_metadata.max_seq_len = min(attn_metadata.max_seq_len,
293 - self.max_model_len)
307 + common_attn_metadata.seq_lens += 1
308 + common_attn_metadata.seq_lens_cpu += 1
309 # For the requests that exceed the max model length, we set the
310 # sequence length to 1 to minimize their overheads in attention.
296 - attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
311 + common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len,
312 + 1)
313 +
314 + common_attn_metadata.num_computed_tokens_cpu = \
315 + common_attn_metadata.seq_lens_cpu - 1
316
317 # Compute the slot mapping.
318 block_numbers = clamped_positions // self.block_size
300 - block_ids = attn_metadata.block_table.gather(
319 + block_ids = common_attn_metadata.block_table_tensor.gather(
320 dim=1, index=block_numbers.view(-1, 1))
321 block_ids = block_ids.view(-1)
303 - attn_metadata.slot_mapping = (block_ids * self.block_size +
304 - clamped_positions % self.block_size)
322 + common_attn_metadata.slot_mapping = (
323 + block_ids * self.block_size +
324 + clamped_positions % self.block_size)
325 # Mask out the slot mappings that exceed the max model length.
326 # Otherwise, the KV cache will be inadvertently updated with the
327 # padding tokens.
308 - attn_metadata.slot_mapping.masked_fill_(exceeds_max_model_len,
309 - PADDING_SLOT_ID)
328 + common_attn_metadata.slot_mapping.masked_fill_(
329 + exceeds_max_model_len, PADDING_SLOT_ID)
330 +
331 + # Rebuild attention metadata
332 + attn_metadata_builder = \
333 + self.runner.attn_groups[0][0].metadata_builders[ubatch_id]
334 + attn_metadata = attn_metadata_builder\
335 + .build_for_drafting(common_attn_metadata=common_attn_metadata,
336 + draft_index=token_index + 1)
337 + for layer_name in self.attn_layer_names:
338 + per_layer_attn_metadata[layer_name] = attn_metadata
339
340 # copy inputs to buffer for cudagraph
341 self.input_ids[:batch_size] = input_ids
@@ -347,6 +376,158 @@ class EagleProposer:
376 draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
377 return draft_token_ids
378
379 + def prepare_next_token_ids_cpu(
380 + self, sampled_token_ids: list[list[int]],
381 + requests: dict[str,
382 + CachedRequestState], gpu_input_batch: InputBatch,
383 + num_scheduled_tokens: dict[str, int]) -> torch.Tensor:
384 + """
385 + This function is used to prepare the inputs for speculative decoding.
386 + It calculates the next token ids for each request based on the sampled
387 + token ids from the CPU. If a request has no sampled token ids (e.g.,
388 + during the initial decoding steps), it falls back to using the request
389 + state to get the next token id.
390 + """
391 + req_ids = gpu_input_batch.req_ids
392 + next_token_ids: list[int] = []
393 + for i, token_ids in enumerate(sampled_token_ids):
394 + if token_ids:
395 + # Common case.
396 + next_token_id = token_ids[-1]
397 + else:
398 + # Partial prefill (rare case).
399 + # Get the next token id from the request state.
400 + req_id = req_ids[i]
401 + req_state = requests[req_id]
402 + seq_len = (req_state.num_computed_tokens +
403 + num_scheduled_tokens[req_id])
404 + next_token_id = req_state.get_token_id(seq_len)
405 + next_token_ids.append(next_token_id)
406 + next_token_ids = torch.tensor(next_token_ids,
407 + dtype=torch.int32,
408 + device=self.input_ids.device)
409 + return next_token_ids
410 +
411 + def prepare_next_token_ids_padded(self,
412 + common_attn_metadata: CommonAttentionMetadata,
413 + sampled_token_ids: torch.Tensor,
414 + requests: dict[str, CachedRequestState],
415 + gpu_input_batch: InputBatch,
416 + discard_request_indices: torch.Tensor,
417 + num_discarded_requests: int) -> \
418 + tuple[torch.Tensor, torch.Tensor]:
419 + """
420 + This function is used to prepare the inputs for speculative decoding.
421 + It calculates the next token ids and the number of valid sampled tokens
422 + for each request, considering the "discarded" requests whose next token
423 + is not sampled and comes from `request.get_token_id()` instead.
424 + It also accounts for the rejected tokens in `sampled_token_ids`.
425 + This function must use device functions to operate on the inputs, and
426 + should not introduce any blocking CPU-GPU synchronization.
427 + """
428 + # TODO(Ben): Combine this into a custom fused kernel
429 +
430 + # Precompute get_token_id for when there is no valid next token
431 + num_reqs = gpu_input_batch.num_reqs
432 + self.backup_next_token_ids.np[:num_reqs] = np.array([
433 + requests[gpu_input_batch.req_ids[i]].get_token_id(
434 + common_attn_metadata.seq_lens_cpu[i].item())
435 + for i in range(num_reqs)
436 + ])
437 + self.backup_next_token_ids.copy_to_gpu(num_reqs)
438 +
439 + # Mask out the sampled tokens indices that should not be sampled.
440 + discard_sampled_tokens_req_indices = \
441 + discard_request_indices[:num_discarded_requests]
442 +
443 + valid_sampled_token_ids_gpu = sampled_token_ids.clone()
444 + valid_sampled_token_ids_gpu.index_fill_(
445 + 0, discard_sampled_tokens_req_indices, -1)
446 +
447 + # Generate a mask for all valid tokens within those requests
448 + max_gen_len = sampled_token_ids.shape[-1]
449 + if max_gen_len == 1:
450 + valid_mask = torch.ones_like(valid_sampled_token_ids_gpu,
451 + dtype=torch.bool)
452 + else:
453 + valid_mask = (
454 + (valid_sampled_token_ids_gpu != -1) &
455 + (valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size))
456 +
457 + # Count the number of valid tokens in each request
458 + valid_sampled_tokens_count = valid_mask.sum(dim=1)
459 +
460 + # Get the rightmost valid index per row
461 + last_valid_indices = valid_sampled_tokens_count - 1
462 + last_valid_indices_safe = torch.clamp(last_valid_indices, min=0)
463 +
464 + # Get last valid token from each row
465 + # (assume undefined state where there is no valid token)
466 + selected_tokens = torch.gather(
467 + valid_sampled_token_ids_gpu, 1,
468 + last_valid_indices_safe.unsqueeze(1)).squeeze(1)
469 +
470 + # Use last token if valid, pre-computed backup if not
471 + batch_size = valid_sampled_token_ids_gpu.shape[0]
472 + next_token_ids = torch.where(
473 + last_valid_indices != -1, selected_tokens,
474 + self.backup_next_token_ids.gpu[:batch_size])
475 +
476 + return next_token_ids, valid_sampled_tokens_count
477 +
478 + def prepare_inputs_padded(self,
479 + common_attn_metadata: CommonAttentionMetadata,
480 + spec_decode_metadata: SpecDecodeMetadata,
481 + valid_sampled_tokens_count: torch.Tensor) -> \
482 + tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
483 + """
484 + This function is used to prepare the inputs for speculative decoding
485 + It updates the common_attn_metadata for speculative decoding,
486 + but does not consider the rejected tokens. Instead, all tokens
487 + are included as inputs to the speculator, with the rejected tokens
488 + used as padding and filtered out later by `token_indices_to_sample`.
489 + No blocking CPU operations should be introduced in this function.
490 + """
491 + num_draft_tokens_gpu = torch.cat([
492 + spec_decode_metadata.cu_num_draft_tokens[0:1],
493 + spec_decode_metadata.cu_num_draft_tokens[1:] -
494 + spec_decode_metadata.cu_num_draft_tokens[:-1]
495 + ])
496 +
497 + num_rejected_tokens_gpu = torch.where(
498 + num_draft_tokens_gpu > 0,
499 + num_draft_tokens_gpu + 1 - valid_sampled_tokens_count,
500 + torch.zeros_like(num_draft_tokens_gpu))
501 +
502 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
503 +
504 + new_query_len_per_req = (query_start_loc_cpu[1:] -
505 + query_start_loc_cpu[:-1])
506 +
507 + total_num_tokens = query_start_loc_cpu[-1].item()
508 + token_indices = self.arange[:total_num_tokens]
509 +
510 + spec_common_attn_metadata = CommonAttentionMetadata(
511 + query_start_loc=common_attn_metadata.query_start_loc,
512 + seq_lens=common_attn_metadata.seq_lens,
513 + query_start_loc_cpu=query_start_loc_cpu,
514 + seq_lens_cpu=common_attn_metadata.seq_lens_cpu,
515 + num_computed_tokens_cpu=common_attn_metadata.
516 + num_computed_tokens_cpu,
517 + num_reqs=common_attn_metadata.num_reqs,
518 + num_actual_tokens=total_num_tokens,
519 + max_query_len=new_query_len_per_req.max().item(),
520 + max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(),
521 + block_table_tensor=common_attn_metadata.block_table_tensor,
522 + slot_mapping=common_attn_metadata.slot_mapping[token_indices],
523 + causal=True,
524 + )
525 +
526 + token_indices_to_sample = common_attn_metadata.query_start_loc[1:] - 1 \
527 + - num_rejected_tokens_gpu
528 +
529 + return spec_common_attn_metadata, token_indices, token_indices_to_sample
530 +
531 def propose_tree(
532 self,
533 batch_size: int,
@@ -520,11 +701,11 @@ class EagleProposer:
701 def prepare_inputs(
702 self,
703 common_attn_metadata: CommonAttentionMetadata,
523 - # [batch_size]
524 - num_rejected_tokens: torch.Tensor
704 + sampled_token_ids: list[list[int]],
705 + num_draft_tokens: list[int],
706 ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
707 """
527 - This function is used to prepare the inputs for the spec decode.
708 + This function is used to prepare the inputs for speculative decoding.
709 It updates to the common_attn_metadata to account for the rejected
710 tokens (and newly sampled tokens). It also returns the token indices
711 of the tokens that should be fed to the speculator.
@@ -545,6 +726,13 @@ class EagleProposer:
726 # q1, q1 + 1, ..., q1 + q2 - n2 - 1,
727 # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
728
729 + num_rejected_tokens = [
730 + n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
731 + for i, n in enumerate(num_draft_tokens)
732 + ]
733 + num_rejected_tokens = torch.tensor(num_rejected_tokens,
734 + dtype=torch.int32)
735 +
736 device = common_attn_metadata.query_start_loc.device
737 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
738 new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu \
vllm/v1/worker/gpu_input_batch.py
+4 -1
@@ -64,7 +64,10 @@ class CachedRequestState:
64 def get_token_id(self, idx: int) -> int:
65 if idx < self.num_prompt_tokens:
66 return self.prompt_token_ids[idx]
67 - return self.output_token_ids[idx - self.num_prompt_tokens]
67 + elif idx - self.num_prompt_tokens < len(self.output_token_ids):
68 + return self.output_token_ids[idx - self.num_prompt_tokens]
69 + else:
70 + return -1
71
72
73 class InputBatch:
vllm/v1/worker/gpu_model_runner.py
+101 -63
@@ -344,6 +344,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
344 self.hidden_size,
345 dtype=self.dtype,
346 numpy=False)
347 + self.discard_request_indices = self._make_buffer(self.max_num_reqs,
348 + dtype=torch.int64)
349 + self.num_discarded_requests = 0
350 +
351 self.num_draft_tokens = self._make_buffer(self.max_num_reqs,
352 dtype=torch.int32)
353 self.num_accepted_tokens = self._make_buffer(self.max_num_reqs,
@@ -974,6 +978,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
978 seq_lens = self.seq_lens.gpu[:num_reqs]
979 max_seq_len = self.seq_lens.np[:num_reqs].max().item()
980
981 + num_tokens = [
982 + self.requests[r].num_tokens for r in self.input_batch.req_ids
983 + ]
984 + num_tokens_np = np.array(num_tokens, dtype=np.int32)
985 +
986 + # Record the index of requests that should not be sampled,
987 + # so that we could clear the sampled tokens before returning
988 + discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np
989 + discard_request_indices = np.nonzero(discard_requests_mask)[0]
990 + self.num_discarded_requests = len(discard_request_indices)
991 + self.discard_request_indices.np[:self.num_discarded_requests] = (
992 + discard_request_indices)
993 +
994 + self.discard_request_indices.copy_to_gpu(self.num_discarded_requests)
995 +
996 # Copy the tensors to the GPU.
997 self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens)
998
@@ -1973,23 +1992,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1992 if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
1993 num_nans_in_logits = self._get_nans_in_logits(logits)
1994
1976 - # TODO(woosuk): The following loop can be slow since it iterates over
1977 - # the requests one by one. Optimize.
1978 - discard_sampled_tokens_req_indices = []
1979 - for i, req_id in enumerate(self.input_batch.req_ids):
1980 - req_state = self.requests[req_id]
1981 - seq_len = (req_state.num_computed_tokens +
1982 - scheduler_output.num_scheduled_tokens[req_id])
1983 - if seq_len < req_state.num_tokens:
1984 - # Ignore the sampled token for partial prefills.
1985 - # Rewind the generator state as if the token was not sampled.
1986 - # This relies on cuda-specific torch-internal impl details
1987 - generator = self.input_batch.generators.get(i)
1988 - if generator is not None:
1989 - generator.set_offset(generator.get_offset() - 4)
1990 - # Record the index of the request that should not be sampled,
1991 - # so that we could clear the sampled tokens before returning.
1992 - discard_sampled_tokens_req_indices.append(i)
1995 + discard_sampled_tokens_req_indices = \
1996 + self.discard_request_indices.np[:self.num_discarded_requests]
1997 + for i in discard_sampled_tokens_req_indices:
1998 + gen = self.input_batch.generators.get(int(i))
1999 + if gen is not None:
2000 + gen.set_offset(gen.get_offset() - 4)
2001
2002 # Copy some objects so they don't get modified after returning.
2003 # This is important when using async scheduling.
@@ -2026,10 +2034,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2034 )
2035 # Mask out the sampled tokens that should not be sampled.
2036 for i in discard_sampled_tokens_req_indices:
2029 - valid_sampled_token_ids[i].clear()
2037 + valid_sampled_token_ids[int(i)].clear()
2038 else:
2039 valid_sampled_token_ids = []
2032 - invalid_req_indices = list(discard_sampled_tokens_req_indices)
2040 + invalid_req_indices = discard_sampled_tokens_req_indices.tolist()
2041 invalid_req_indices_set = set(invalid_req_indices)
2042 assert sampled_token_ids.shape[-1] == 1
2043
@@ -2229,6 +2237,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2237 with record_function_or_nullcontext("Sample"):
2238 sampler_output = self._sample(logits, spec_decode_metadata)
2239
2240 + def propose_draft_token_ids(sampled_token_ids):
2241 + assert spec_decode_common_attn_metadata is not None
2242 + with record_function_or_nullcontext("Draft"):
2243 + self._draft_token_ids = self.propose_draft_token_ids(
2244 + scheduler_output,
2245 + sampled_token_ids,
2246 + self.input_batch.sampling_metadata,
2247 + hidden_states,
2248 + sample_hidden_states,
2249 + aux_hidden_states,
2250 + spec_decode_metadata,
2251 + spec_decode_common_attn_metadata,
2252 + )
2253 +
2254 + use_padded_batch_for_eagle = self.speculative_config and \
2255 + self.speculative_config.use_eagle() and \
2256 + not self.speculative_config.disable_padded_drafter_batch
2257 + if use_padded_batch_for_eagle:
2258 + # EAGLE speculative decoding can use the GPU sampled tokens
2259 + # as inputs, and does not need to wait for bookkeeping to finish.
2260 + propose_draft_token_ids(sampler_output.sampled_token_ids)
2261 +
2262 with record_function_or_nullcontext("Bookkeep"):
2263 (
2264 num_nans_in_logits,
@@ -2242,19 +2272,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2272 logits, hidden_states,
2273 num_scheduled_tokens)
2274
2245 - if self.speculative_config:
2246 - assert spec_decode_common_attn_metadata is not None
2247 - with record_function_or_nullcontext("Draft"):
2248 - self._draft_token_ids = self.propose_draft_token_ids(
2249 - scheduler_output,
2250 - valid_sampled_token_ids,
2251 - self.input_batch.sampling_metadata,
2252 - hidden_states,
2253 - sample_hidden_states,
2254 - aux_hidden_states,
2255 - spec_decode_metadata,
2256 - spec_decode_common_attn_metadata,
2257 - )
2275 + if self.speculative_config and not use_padded_batch_for_eagle:
2276 + # ngram and other speculative decoding methods use the sampled
2277 + # tokens on the CPU, so they are run after bookkeeping.
2278 + propose_draft_token_ids(valid_sampled_token_ids)
2279
2280 with record_function_or_nullcontext("EPLB"):
2281 self.eplb_step()
@@ -2294,7 +2315,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2315 def propose_draft_token_ids(
2316 self,
2317 scheduler_output: "SchedulerOutput",
2297 - sampled_token_ids: list[list[int]],
2318 + sampled_token_ids: Union[torch.Tensor, list[list[int]]],
2319 sampling_metadata: SamplingMetadata,
2320 hidden_states: torch.Tensor,
2321 sample_hidden_states: torch.Tensor,
@@ -2304,11 +2325,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2325 ) -> Union[list[list[int]], torch.Tensor]:
2326 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2327 if self.speculative_config.method == "ngram":
2328 + assert isinstance(sampled_token_ids, list)
2329 assert isinstance(self.drafter, NgramProposer)
2330 draft_token_ids = self.propose_ngram_draft_token_ids(
2331 sampled_token_ids)
2332 elif self.speculative_config.method == "medusa":
2333 + assert isinstance(sampled_token_ids, list)
2334 assert isinstance(self.drafter, MedusaProposer)
2335 +
2336 if sample_hidden_states.shape[0] == len(sampled_token_ids):
2337 # The input to the target model does not include draft tokens.
2338 hidden_states = sample_hidden_states
@@ -2329,27 +2353,37 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2353 )
2354 elif self.speculative_config.use_eagle():
2355 assert isinstance(self.drafter, EagleProposer)
2332 - # TODO(woosuk): Refactor the loop.
2333 - req_ids = self.input_batch.req_ids
2334 - next_token_ids: list[int] = []
2335 - for i, token_ids in enumerate(sampled_token_ids):
2336 - if token_ids:
2337 - # Common case.
2338 - next_token_id = token_ids[-1]
2339 - else:
2340 - # Partial prefill (rare case).
2341 - # Get the next token id from the request state.
2342 - req_id = req_ids[i]
2343 - req_state = self.requests[req_id]
2344 - seq_len = (req_state.num_computed_tokens +
2345 - scheduler_output.num_scheduled_tokens[req_id])
2346 - next_token_id = req_state.get_token_id(seq_len)
2347 - next_token_ids.append(next_token_id)
2348 - next_token_ids = torch.tensor(next_token_ids,
2349 - dtype=torch.int32,
2350 - device=self.device)
2356 +
2357 + if self.speculative_config.disable_padded_drafter_batch:
2358 + # When padded-batch is disabled, the sampled_token_ids should be
2359 + # the cpu-side list[list[int]] of valid sampled tokens for each
2360 + # request, with invalid requests having empty lists.
2361 + assert isinstance(sampled_token_ids, list), \
2362 + "sampled_token_ids should be a python list when" \
2363 + "padded-batch is disabled."
2364 + next_token_ids = self.drafter.prepare_next_token_ids_cpu(
2365 + sampled_token_ids, self.requests, self.input_batch,
2366 + scheduler_output.num_scheduled_tokens)
2367 + else:
2368 + # When using padded-batch, the sampled_token_ids should be
2369 + # the gpu tensor of sampled tokens for each request, of shape
2370 + # (num_reqs, num_spec_tokens + 1) with rejected tokens having
2371 + # value -1.
2372 + assert isinstance(sampled_token_ids, torch.Tensor), \
2373 + "sampled_token_ids should be a torch.Tensor when" \
2374 + "padded-batch is enabled."
2375 + next_token_ids, valid_sampled_tokens_count = \
2376 + self.drafter.prepare_next_token_ids_padded(
2377 + common_attn_metadata,
2378 + sampled_token_ids,
2379 + self.requests,
2380 + self.input_batch,
2381 + self.discard_request_indices.gpu,
2382 + self.num_discarded_requests
2383 + )
2384
2385 if spec_decode_metadata is None:
2386 + token_indices_to_sample = None
2387 # input_ids can be None for multimodal models.
2388 target_token_ids = self.input_ids.gpu[:num_scheduled_tokens]
2389 # TODO(woosuk): Support M-RoPE.
@@ -2361,17 +2395,20 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2395 else:
2396 target_hidden_states = hidden_states[:num_scheduled_tokens]
2397 else:
2364 - # TODO(woosuk): Refactor this.
2365 - num_draft_tokens = spec_decode_metadata.num_draft_tokens
2366 - num_rejected_tokens = [
2367 - n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
2368 - for i, n in enumerate(num_draft_tokens)
2369 - ]
2370 - num_rejected_tokens_cpu = torch.tensor(num_rejected_tokens,
2371 - dtype=torch.int32)
2372 - common_attn_metadata, token_indices =\
2373 - self.drafter.prepare_inputs(
2374 - common_attn_metadata, num_rejected_tokens_cpu)
2398 + if self.speculative_config.disable_padded_drafter_batch:
2399 + token_indices_to_sample = None
2400 + common_attn_metadata, token_indices =\
2401 + self.drafter.prepare_inputs(
2402 + common_attn_metadata,
2403 + sampled_token_ids,
2404 + spec_decode_metadata.num_draft_tokens)
2405 + else:
2406 + common_attn_metadata, token_indices, \
2407 + token_indices_to_sample =\
2408 + self.drafter.prepare_inputs_padded(
2409 + common_attn_metadata,
2410 + spec_decode_metadata,
2411 + valid_sampled_tokens_count)
2412
2413 target_token_ids = self.input_ids.gpu[token_indices]
2414 # TODO(woosuk): Support M-RoPE.
@@ -2391,6 +2428,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2428 target_positions=target_positions,
2429 target_hidden_states=target_hidden_states,
2430 next_token_ids=next_token_ids,
2431 + last_token_indices=token_indices_to_sample,
2432 sampling_metadata=sampling_metadata,
2433 common_attn_metadata=common_attn_metadata,
2434 mm_embeds=mm_embeds,