@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b004c00418
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 12, 16:00 UTC

[Model Runner V2] Support spec decoding [1/N] (#29274)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Nov 23, 2025 at 10:09 UTC b004c00418268daa61b3526358b661165a360f7d
5 files changed +350 -29
vllm/v1/worker/gpu/input_batch.py
+112 -16
@@ -35,6 +35,7 @@ class InputBuffers:
35 self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
36 self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
37 self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
38 + self.cu_num_logits = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
39
40 # Structured outputs.
41 self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
@@ -64,6 +65,7 @@ class InputBatch:
65 # sum(num_scheduled_tokens)
66 num_tokens: int
67 num_tokens_after_padding: int
68 + num_draft_tokens: int
69
70 # [num_reqs + 1]
71 query_start_loc: torch.Tensor
@@ -80,8 +82,10 @@ class InputBatch:
82 # layer_name -> Metadata
83 attn_metadata: dict[str, Any]
84
83 - # [num_reqs]
85 + # [total_num_logits]
86 logits_indices: torch.Tensor
87 + # [num_reqs + 1]
88 + cu_num_logits: torch.Tensor
89
90 @classmethod
91 def make_dummy(
@@ -118,6 +122,7 @@ class InputBatch:
122 positions = input_buffers.positions[:num_tokens]
123 # attn_metadata = defaultdict(lambda: None)
124 logits_indices = query_start_loc[1:] - 1
125 + cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
126 return cls(
127 req_ids=req_ids,
128 num_reqs=num_reqs,
@@ -126,6 +131,7 @@ class InputBatch:
131 num_scheduled_tokens=num_scheduled_tokens,
132 num_tokens=num_tokens,
133 num_tokens_after_padding=num_tokens,
134 + num_draft_tokens=0,
135 query_start_loc=query_start_loc,
136 query_start_loc_np=query_start_loc_np,
137 seq_lens=seq_lens,
@@ -134,6 +140,7 @@ class InputBatch:
140 positions=positions,
141 attn_metadata=None, # type: ignore
142 logits_indices=logits_indices,
143 + cu_num_logits=cu_num_logits,
144 )
145
146
@@ -279,19 +286,53 @@ def _combine_sampled_and_draft_tokens_kernel(
286 query_start_loc_ptr,
287 seq_lens_ptr,
288 prefill_len_ptr,
289 + draft_tokens_ptr,
290 + draft_tokens_stride,
291 + cu_num_logits_ptr,
292 + logits_indices_ptr,
293 + BLOCK_SIZE: tl.constexpr,
294 ):
295 batch_idx = tl.program_id(0)
296 req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
297
298 + # Get the number of logits and draft tokens.
299 + cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx)
300 + cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
301 + num_logits = cu_num_logits_end - cu_num_logits_start
302 + num_draft_tokens = num_logits - 1
303 +
304 + # Compute the logits indices.
305 + block = tl.arange(0, BLOCK_SIZE)
306 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
307 + logits_start = query_end - num_logits
308 + tl.store(
309 + logits_indices_ptr + cu_num_logits_start + block,
310 + logits_start + block,
311 + mask=block < num_logits,
312 + )
313 +
314 seq_len = tl.load(seq_lens_ptr + batch_idx)
315 prefill_len = tl.load(prefill_len_ptr + req_state_idx)
316 if seq_len <= prefill_len:
289 - # Handling prefill tokens.
317 + # Handling prefill tokens. No sampled or draft tokens.
318 return
319
320 + # Write the last sampled token ID to input_ids.
321 last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
293 - end = tl.load(query_start_loc_ptr + batch_idx + 1)
294 - tl.store(input_ids_ptr + end - 1, last_token_id)
322 + tl.store(input_ids_ptr + query_end - num_logits, last_token_id)
323 +
324 + # Write the draft tokens (if any) to input_ids.
325 + if num_draft_tokens > 0:
326 + mask = block < num_draft_tokens
327 + draft_tokens = tl.load(
328 + draft_tokens_ptr + req_state_idx * draft_tokens_stride + block,
329 + mask=mask,
330 + )
331 + tl.store(
332 + input_ids_ptr + query_end - num_draft_tokens + block,
333 + draft_tokens,
334 + mask=mask,
335 + )
336
337
338 def combine_sampled_and_draft_tokens(
@@ -301,8 +342,18 @@ def combine_sampled_and_draft_tokens(
342 query_start_loc: torch.Tensor,
343 seq_lens: torch.Tensor,
344 prefill_len: torch.Tensor,
345 + draft_tokens: torch.Tensor,
346 + cu_num_logits: torch.Tensor,
347 + num_logits: int,
348 ) -> torch.Tensor:
349 num_reqs = seq_lens.shape[0]
350 + num_speculative_steps = draft_tokens.shape[-1]
351 +
352 + logits_indices = torch.empty(
353 + num_logits,
354 + dtype=torch.int64,
355 + device=input_ids.device,
356 + )
357 _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
358 input_ids,
359 idx_mapping,
@@ -310,35 +361,80 @@ def combine_sampled_and_draft_tokens(
361 query_start_loc,
362 seq_lens,
363 prefill_len,
364 + draft_tokens,
365 + draft_tokens.stride(0),
366 + cu_num_logits,
367 + logits_indices,
368 + # NOTE(woosuk): Add 1 to ensure the block can cover the last sampled token
369 + # in addition to all draft tokens.
370 + BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
371 )
314 - return input_ids
372 + return logits_indices
373
374
375 @triton.jit
318 -def _update_num_computed_tokens_kernel(
376 +def _post_update_kernel(
377 idx_mapping_ptr,
378 num_computed_tokens_ptr,
379 + last_sampled_tokens_ptr,
380 + sampled_tokens_ptr,
381 + sampled_tokens_stride,
382 + num_sampled_ptr,
383 query_start_loc_ptr,
384 + cu_num_logits_ptr,
385 ):
386 req_id = tl.program_id(0)
387 req_state_idx = tl.load(idx_mapping_ptr + req_id)
388
326 - start = tl.load(query_start_loc_ptr + req_id)
327 - end = tl.load(query_start_loc_ptr + req_id + 1)
328 - query_len = end - start
329 -
330 - n = tl.load(num_computed_tokens_ptr + req_state_idx)
331 - tl.store(num_computed_tokens_ptr + req_state_idx, n + query_len)
332 -
333 -
334 -def update_num_computed_tokens(
389 + num_sampled = tl.load(num_sampled_ptr + req_id)
390 + if num_sampled > 0:
391 + token_id = tl.load(
392 + sampled_tokens_ptr + req_id * sampled_tokens_stride + num_sampled - 1
393 + )
394 + tl.store(last_sampled_tokens_ptr + req_state_idx, token_id)
395 +
396 + query_start = tl.load(query_start_loc_ptr + req_id)
397 + query_end = tl.load(query_start_loc_ptr + req_id + 1)
398 + query_len = query_end - query_start
399 +
400 + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
401 + num_computed += query_len
402 + # Consider the rejected tokens in spec decoding.
403 + if num_sampled > 0:
404 + # NOTE(woosuk): We must skip num_sampled == 0 to account for chunked prefills.
405 + logits_start = tl.load(cu_num_logits_ptr + req_id)
406 + logits_end = tl.load(cu_num_logits_ptr + req_id + 1)
407 + num_logits = logits_end - logits_start
408 + num_rejected = num_logits - num_sampled
409 + num_computed -= num_rejected
410 + tl.store(num_computed_tokens_ptr + req_state_idx, num_computed)
411 +
412 +
413 +def post_update(
414 + # [num_reqs]
415 idx_mapping: torch.Tensor,
416 + # [max_num_reqs]
417 num_computed_tokens: torch.Tensor,
418 + # [max_num_reqs]
419 + last_sampled_tokens: torch.Tensor,
420 + # [num_reqs, num_speculative_steps + 1]
421 + sampled_tokens: torch.Tensor,
422 + # [num_reqs]
423 + num_sampled: torch.Tensor,
424 + # [num_reqs + 1]
425 query_start_loc: torch.Tensor,
426 + # [num_reqs + 1]
427 + cu_num_logits: torch.Tensor,
428 ) -> None:
429 num_reqs = idx_mapping.shape[0]
340 - _update_num_computed_tokens_kernel[(num_reqs,)](
430 + _post_update_kernel[(num_reqs,)](
431 idx_mapping,
432 num_computed_tokens,
433 + last_sampled_tokens,
434 + sampled_tokens,
435 + sampled_tokens.stride(0),
436 + num_sampled,
437 query_start_loc,
438 + cu_num_logits,
439 + num_warps=1,
440 )
vllm/v1/worker/gpu/model_runner.py
+73 -13
@@ -40,11 +40,12 @@ from vllm.v1.worker.gpu.input_batch import (
40 InputBatch,
41 InputBuffers,
42 combine_sampled_and_draft_tokens,
43 + post_update,
44 prepare_pos_seq_lens,
45 prepare_prefill_inputs,
45 - update_num_computed_tokens,
46 )
47 from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
48 +from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
49 from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
50 from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
51 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -100,10 +101,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
101 self.input_prep_event = None
102 self.structured_outputs_event = None
103
104 + if self.speculative_config is not None:
105 + self.do_spec_decode = True
106 + self.num_speculative_steps = self.speculative_config.num_speculative_tokens
107 + else:
108 + self.do_spec_decode = False
109 + self.num_speculative_steps = 0
110 +
111 self.req_states = RequestState(
112 max_num_reqs=self.max_num_reqs,
113 max_model_len=self.max_model_len,
114 max_num_batched_tokens=self.max_num_tokens,
115 + num_speculative_steps=self.num_speculative_steps,
116 vocab_size=self.vocab_size,
117 device=self.device,
118 pin_memory=self.pin_memory,
@@ -427,6 +436,32 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
436 idx_mapping_np = idx_mapping.np[:num_reqs]
437 idx_mapping = idx_mapping.copy_to_gpu(num_reqs)
438
439 + # Get the number of draft tokens for each request.
440 + if not scheduler_output.scheduled_spec_decode_tokens:
441 + # No draft token scheduled (common case).
442 + total_num_draft_tokens = 0
443 + total_num_logits = num_reqs
444 + cu_num_logits = torch.arange(
445 + num_reqs + 1, device=self.device, dtype=torch.int32
446 + )
447 + else:
448 + draft_tokens = scheduler_output.scheduled_spec_decode_tokens
449 + num_draft_tokens = np.array(
450 + [
451 + len(draft_tokens[req_id]) if req_id in draft_tokens else 0
452 + for req_id in req_ids
453 + ],
454 + dtype=np.int32,
455 + )
456 + total_num_draft_tokens = int(num_draft_tokens.sum())
457 + total_num_logits = num_reqs + total_num_draft_tokens
458 +
459 + np.cumsum(
460 + num_draft_tokens + 1,
461 + out=self.input_buffers.cu_num_logits.np[1 : num_reqs + 1],
462 + )
463 + cu_num_logits = self.input_buffers.cu_num_logits.copy_to_gpu(num_reqs + 1)
464 +
465 # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
466 block_tables = self.block_tables.gather_block_tables(idx_mapping)
467
@@ -456,14 +491,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
491 seq_lens = self.input_buffers.seq_lens[:num_reqs]
492
493 # Some input token ids are directly read from the last sampled tokens
459 - # and draft tokens.
460 - combine_sampled_and_draft_tokens(
494 + # and draft tokens. Also, get the logits indices to sample tokens from.
495 + logits_indices = combine_sampled_and_draft_tokens(
496 self.input_buffers.input_ids.gpu,
497 idx_mapping,
498 self.req_states.last_sampled_tokens,
499 query_start_loc_gpu,
500 seq_lens,
501 self.req_states.prefill_len.gpu,
502 + self.req_states.draft_tokens,
503 + cu_num_logits,
504 + total_num_logits,
505 )
506
507 # Compute slot mappings: [num_kv_cache_groups, num_tokens]
@@ -471,9 +509,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
509 query_start_loc_gpu, self.input_buffers.positions[:num_tokens]
510 )
511
474 - # Logits indices to sample next token from.
475 - logits_indices = query_start_loc_gpu[1:] - 1
476 -
512 # Get num_computed_tokens.
513 # HACK(woosuk): Here, we use num_computed_tokens on GPU instead of
514 # num_computed_tokens_cpu. This works for most cases.
@@ -508,6 +543,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
543 num_scheduled_tokens=num_scheduled_tokens,
544 num_tokens=num_tokens,
545 num_tokens_after_padding=num_tokens_after_padding,
546 + num_draft_tokens=total_num_draft_tokens,
547 query_start_loc=query_start_loc_gpu,
548 query_start_loc_np=query_start_loc_np,
549 seq_lens=seq_lens,
@@ -516,6 +552,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
552 positions=positions,
553 attn_metadata=attn_metadata,
554 logits_indices=logits_indices,
555 + cu_num_logits=cu_num_logits,
556 )
557
558 def sample(
@@ -530,6 +567,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
567 if grammar_output is not None:
568 # Apply grammar bitmask to the logits in-place.
569 # TODO(woosuk): Make compatible with spec decoding.
570 + assert input_batch.num_draft_tokens == 0
571 with async_barrier(self.structured_outputs_event):
572 apply_grammar_bitmask(
573 logits,
@@ -539,12 +577,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
577 self.input_buffers,
578 )
579
580 + # Sample tokens and compute logprobs (if needed).
581 sampler_output = self.sampler(logits, sampling_metadata)
582 +
583 # Get the number of sampled tokens.
544 - # 0 if chunked-prefilling, 1 if not.
584 prefill_len = self.req_states.prefill_len.gpu[input_batch.idx_mapping]
585 is_chunked_prefilling = input_batch.seq_lens < prefill_len
547 - num_sampled = (~is_chunked_prefilling).int()
586 + if input_batch.num_draft_tokens == 0:
587 + # No draft tokens (common case).
588 + # 0 if chunked-prefilling, 1 if not.
589 + num_sampled = (~is_chunked_prefilling).int()
590 + else:
591 + # Draft tokens for spec decoding.
592 + input_ids = input_batch.input_ids[input_batch.logits_indices]
593 + sampled_tokens, num_sampled = rejection_sample(
594 + sampler_output.sampled_token_ids,
595 + input_ids,
596 + input_batch.cu_num_logits,
597 + self.num_speculative_steps,
598 + )
599 + num_sampled *= ~is_chunked_prefilling
600 + sampler_output.sampled_token_ids = sampled_tokens
601 + # TODO(woosuk): Support logprobs with spec decoding.
602 return sampler_output, num_sampled
603
604 def compute_prompt_logprobs(
@@ -653,11 +707,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
707 num_sampled: torch.Tensor,
708 ) -> None:
709 # Update the number of computed tokens.
656 - update_num_computed_tokens(
710 + post_update(
711 input_batch.idx_mapping,
712 self.req_states.num_computed_tokens,
713 + self.req_states.last_sampled_tokens,
714 + sampled_tokens,
715 + num_sampled,
716 input_batch.query_start_loc,
717 + input_batch.cu_num_logits,
718 )
719 +
720 + # Update the number of computed prefill tokens.
721 idx_mapping_np = input_batch.idx_mapping_np
722 computed_prefill = self.req_states.num_computed_prefill_tokens
723 # TODO(woosuk): Simplify this.
@@ -666,10 +726,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
726 self.req_states.prefill_len.np[idx_mapping_np],
727 )
728
669 - # Store the last sampled token ids.
670 - last_sampled = sampled_tokens
671 - self.req_states.last_sampled_tokens[input_batch.idx_mapping] = last_sampled
672 -
729 def get_cudagraph_and_dp_padding(
730 self,
731 scheduler_output: SchedulerOutput,
@@ -761,6 +817,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
817 sampling_metadata = self.req_states.make_sampling_metadata(
818 input_batch.idx_mapping_np, pos
819 )
820 + if input_batch.num_draft_tokens > 0:
821 + sampling_metadata = self.req_states.expand_sampling_metadata(
822 + sampling_metadata, input_batch.cu_num_logits
823 + )
824
825 if self.lora_config:
826 # Activate LoRA adapters.
vllm/v1/worker/gpu/spec_decode/__init__.py
vllm/v1/worker/gpu/spec_decode/rejection_sample.py new
+71
@@ -0,0 +1,71 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.triton_utils import tl, triton
6 +
7 +
8 +@triton.jit
9 +def _rejection_sample_kernel(
10 + sampled_ptr, # [num_reqs, num_speculative_steps + 1]
11 + sampled_stride,
12 + num_sampled_ptr, # [num_reqs]
13 + target_sampled_ptr, # [num_draft_tokens + num_reqs]
14 + input_ids_ptr, # [num_draft_tokens + num_reqs]
15 + cu_num_logits_ptr, # [num_reqs + 1]
16 +):
17 + req_idx = tl.program_id(0)
18 + start_idx = tl.load(cu_num_logits_ptr + req_idx)
19 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
20 + num_tokens = end_idx - start_idx
21 +
22 + num_sampled = 0
23 + rejected = False
24 + for i in range(num_tokens - 1):
25 + if not rejected:
26 + target_sampled = tl.load(target_sampled_ptr + start_idx + i)
27 + draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1)
28 + tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled)
29 + num_sampled += 1
30 + if target_sampled != draft_sampled:
31 + rejected = True
32 + if not rejected:
33 + target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
34 + tl.store(
35 + sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled
36 + )
37 + num_sampled += 1
38 + tl.store(num_sampled_ptr + req_idx, num_sampled)
39 +
40 +
41 +def rejection_sample(
42 + # [num_draft_tokens + num_reqs]
43 + target_sampled: torch.Tensor,
44 + # [num_draft_tokens + num_reqs]
45 + input_ids: torch.Tensor,
46 + # [num_reqs + 1]
47 + cu_num_logits: torch.Tensor,
48 + num_speculative_steps: int,
49 +) -> tuple[torch.Tensor, torch.Tensor]:
50 + num_reqs = cu_num_logits.shape[0] - 1
51 + sampled = torch.empty(
52 + num_reqs,
53 + num_speculative_steps + 1,
54 + dtype=target_sampled.dtype,
55 + device=target_sampled.device,
56 + )
57 + num_sampled = torch.empty(
58 + num_reqs,
59 + dtype=torch.int32,
60 + device=target_sampled.device,
61 + )
62 + _rejection_sample_kernel[(num_reqs,)](
63 + sampled,
64 + sampled.stride(0),
65 + num_sampled,
66 + target_sampled,
67 + input_ids,
68 + cu_num_logits,
69 + num_warps=1,
70 + )
71 + return sampled, num_sampled
vllm/v1/worker/gpu/states.py
+94
@@ -7,6 +7,7 @@ import torch
7
8 from vllm.lora.request import LoRARequest
9 from vllm.sampling_params import SamplingParams
10 +from vllm.triton_utils import tl, triton
11 from vllm.v1.outputs import LogprobsTensors
12 from vllm.v1.utils import CpuGpuBuffer
13
@@ -63,6 +64,7 @@ class RequestState:
64 max_num_reqs: int,
65 max_model_len: int,
66 max_num_batched_tokens: int,
67 + num_speculative_steps: int,
68 vocab_size: int,
69 device: torch.device,
70 pin_memory: bool,
@@ -70,6 +72,7 @@ class RequestState:
72 self.max_num_reqs = max_num_reqs
73 self.max_model_len = max_model_len
74 self.max_num_batched_tokens = max_num_batched_tokens
75 + self.num_speculative_steps = num_speculative_steps
76 self.vocab_size = vocab_size
77 self.device = device
78 self.pin_memory = pin_memory
@@ -100,6 +103,14 @@ class RequestState:
103 device=device,
104 )
105
106 + # Draft tokens.
107 + self.draft_tokens = torch.zeros(
108 + self.max_num_reqs,
109 + self.num_speculative_steps,
110 + dtype=torch.int64,
111 + device=device,
112 + )
113 +
114 # LoRA.
115 self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
116 self.lora_ids.fill(NO_LORA_ID)
@@ -226,6 +237,17 @@ class RequestState:
237 max_num_logprobs=max_num_logprobs,
238 )
239
240 + def expand_sampling_metadata(
241 + self,
242 + sampling_metadata: SamplingMetadata,
243 + cu_num_logits: torch.Tensor,
244 + ) -> SamplingMetadata:
245 + # For draft tokens, we need to expand the sampling param tensors as
246 + # each request samples multiple tokens in each step.
247 + return expand_sampling_metadata(
248 + sampling_metadata, cu_num_logits, self.num_speculative_steps
249 + )
250 +
251 def make_lora_inputs(
252 self,
253 req_ids: list[str],
@@ -270,3 +292,75 @@ class Param:
292 class ExtraData:
293 lora_request: LoRARequest | None
294 in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)
295 +
296 +
297 +# NOTE(woosuk): Re-compilation can happen at runtime since top_p and top_k can be None.
298 +@triton.jit
299 +def _expand_sampling_metadata_kernel(
300 + temp_ptr,
301 + expanded_temp_ptr,
302 + top_p_ptr,
303 + expanded_top_p_ptr,
304 + top_k_ptr,
305 + expanded_top_k_ptr,
306 + seeds_ptr,
307 + expanded_seeds_ptr,
308 + cu_num_logits_ptr,
309 + BLOCK_SIZE: tl.constexpr,
310 +):
311 + req_idx = tl.program_id(0)
312 + start_idx = tl.load(cu_num_logits_ptr + req_idx)
313 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
314 + num_tokens = end_idx - start_idx
315 +
316 + block = tl.arange(0, BLOCK_SIZE)
317 + mask = block < num_tokens
318 +
319 + temp = tl.load(temp_ptr + req_idx)
320 + tl.store(expanded_temp_ptr + start_idx + block, temp, mask=mask)
321 +
322 + if top_p_ptr is not None:
323 + top_p = tl.load(top_p_ptr + req_idx)
324 + tl.store(expanded_top_p_ptr + start_idx + block, top_p, mask=mask)
325 +
326 + if top_k_ptr is not None:
327 + top_k = tl.load(top_k_ptr + req_idx)
328 + tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)
329 +
330 + seed = tl.load(seeds_ptr + req_idx)
331 + tl.store(expanded_seeds_ptr + start_idx + block, seed, mask=mask)
332 +
333 +
334 +def expand_sampling_metadata(
335 + sampling_metadata: SamplingMetadata,
336 + cu_num_logits: torch.Tensor,
337 + num_speculative_steps: int,
338 +) -> SamplingMetadata:
339 + total_num_logits = sampling_metadata.pos.shape[0]
340 + create_empty = lambda x: x.new_empty(total_num_logits) if x is not None else None
341 + expanded_temp = create_empty(sampling_metadata.temperature)
342 + expanded_top_p = create_empty(sampling_metadata.top_p)
343 + expanded_top_k = create_empty(sampling_metadata.top_k)
344 + expanded_seeds = create_empty(sampling_metadata.seeds)
345 +
346 + num_reqs = cu_num_logits.shape[0] - 1
347 + _expand_sampling_metadata_kernel[(num_reqs,)](
348 + sampling_metadata.temperature,
349 + expanded_temp,
350 + sampling_metadata.top_p,
351 + expanded_top_p,
352 + sampling_metadata.top_k,
353 + expanded_top_k,
354 + sampling_metadata.seeds,
355 + expanded_seeds,
356 + cu_num_logits,
357 + BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
358 + )
359 + return SamplingMetadata(
360 + temperature=expanded_temp,
361 + top_p=expanded_top_p,
362 + top_k=expanded_top_k,
363 + seeds=expanded_seeds,
364 + pos=sampling_metadata.pos,
365 + max_num_logprobs=sampling_metadata.max_num_logprobs,
366 + )