@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c77181e534

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 6, 20:30 UTC

[Model Runner V2] Add probabilistic rejection sampling for spec decoding (#35461)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed Mar 11, 2026 at 14:04 UTC c77181e534597f7347fc03b7d26600fb3cea9981

9 files changed +494 -112

vllm/config/speculative.py

+10

@@ -57,6 +57,10 @@ SpeculativeMethod = Literal[

57 EagleModelTypes,

58 NgramGPUTypes,

59 ]

60 +RejectionSampleMethod = Literal[

61 + "strict",

62 + "probabilistic",

63 +]

66 @config

@@ -171,6 +175,12 @@ class SpeculativeConfig:

175 """Load config for the draft model. If not specified, will use the load

176 config from the target model."""

177

178 + rejection_sample_method: RejectionSampleMethod = "strict"

179 + """Whether to use strict (target and draft sampled tokens match exactly)

180 + or probabilistic rejection sampling. Both respect the target model

181 + distribution, but the latter yields a higher acceptance rate at the cost

182 + of more memory to cache draft logits."""

183 +

184 def compute_hash(self) -> str:

185 """

186 WARNING: Whenever a new field is added to this config,

vllm/v1/worker/gpu/model_runner.py

+30 -34

@@ -90,7 +90,7 @@ from vllm.v1.worker.gpu.spec_decode import init_speculator

90 from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (

91 set_eagle3_aux_hidden_state_layers,

92 )

93 -from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample

93 +from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler

94 from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler

95 from vllm.v1.worker.gpu.states import RequestState

96 from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker

@@ -162,6 +162,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

162 self.speculator = None

163 self.num_speculative_steps = 0

164 self.use_aux_hidden_state_outputs = False

165 + use_strict_rejection_sampling = False

166 if self.speculative_config is not None:

167 self.num_speculative_steps = self.speculative_config.num_speculative_tokens

168 if self.is_last_pp_rank:

@@ -172,6 +173,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

173 self.use_aux_hidden_state_outputs = True

174 if self.pp_size > 1:

175 raise ValueError("EAGLE3 with pipeline parallel is not supported.")

176 + use_strict_rejection_sampling = (

177 + self.speculative_config.rejection_sample_method == "strict"

178 + )

179

180 # Draft tokens propagation - for spec-dec + struct outputs.

181 self.draft_tokens_handler = DraftTokensHandler(self.device)

@@ -183,6 +187,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

187 num_speculative_steps=self.num_speculative_steps,

188 vocab_size=self.vocab_size,

189 device=self.device,

190 + model_dtype=self.dtype,

191 + cache_draft_logits=not use_strict_rejection_sampling,

192 )

193 self.input_buffers = InputBuffers(

194 max_num_reqs=self.max_num_reqs,

@@ -197,6 +203,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

203 logprobs_mode=self.model_config.logprobs_mode,

204 num_speculative_tokens=self.num_speculative_steps + 1,

205 )

206 + self.rejection_sampler = RejectionSampler(

207 + self.sampler,

208 + num_speculative_steps=self.num_speculative_steps,

209 + use_strict_rejection_sampling=use_strict_rejection_sampling,

210 + )

211 self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)

212

213 # CUDA graphs.

@@ -412,6 +423,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

423 next_prefill_tokens=self.req_states.next_prefill_tokens,

424 temperature=self.sampler.sampling_states.temperature.gpu,

425 seeds=self.sampler.sampling_states.seeds.gpu,

426 + draft_logits_out=self.req_states.draft_logits,

427 num_tokens_across_dp=num_tokens_across_dp,

428 dummy_run=True,

429 skip_attn_for_dummy_run=skip_attn,

@@ -425,24 +437,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

437 def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:

438 num_reqs = hidden_states.shape[0]

439 logits = self.model.compute_logits(hidden_states)

428 - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device)

429 - idx_mapping_np = np.arange(num_reqs, dtype=np.int32)

430 - pos = torch.zeros(num_reqs, dtype=torch.int64, device=self.device)

431 - dummy_input_ids = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)

432 - expanded_local_pos = torch.zeros(

433 - num_reqs, dtype=torch.int32, device=self.device

440 + dummy_input_batch = InputBatch.make_dummy(

441 + num_reqs, num_reqs, self.input_buffers

442 )

443 +

444 # NOTE(woosuk): During the initial memory profiling, the sampler may skip

445 # top_k, top_p, and logprobs, using less GPU memory than what is possible

446 # during actual execution.

447 self.sampler(

448 logits,

440 - idx_mapping,

441 - idx_mapping_np,

442 - idx_mapping_np,

443 - pos,

444 - dummy_input_ids,

445 - expanded_local_pos,

449 + dummy_input_batch,

450 )

451

452 @torch.inference_mode()

@@ -768,8 +772,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

772 grammar_output: GrammarOutput | None,

773 ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:

774 sample_hidden_states = hidden_states[input_batch.logits_indices]

771 - sample_pos = input_batch.positions[input_batch.logits_indices]

772 - input_ids = input_batch.input_ids[input_batch.logits_indices]

775 logits = self.model.compute_logits(sample_hidden_states)

776 if grammar_output is not None:

777 # Apply grammar bitmask to the logits in-place.

@@ -780,34 +782,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):

782 grammar_output.grammar_bitmask,

783 )

784

783 - # Sample tokens and compute logprobs (if needed).

784 - sampler_output = self.sampler(

785 - logits,

786 - input_batch.expanded_idx_mapping,

787 - input_batch.idx_mapping_np,

788 - input_batch.cu_num_logits_np,

789 - sample_pos,

790 - input_ids,

791 - input_batch.expanded_local_pos,

792 - )

793 -

785 if input_batch.num_draft_tokens == 0:

786 # No draft tokens (common case).

796 - num_sampled = input_batch.seq_lens.new_ones(input_batch.num_reqs)

787 + sampler_output = self.sampler(

788 + logits,

789 + input_batch,

790 + )

791 else:

792 # Rejection sampling for spec decoding.

799 - sampled_tokens, num_sampled = rejection_sample(

800 - sampler_output.sampled_token_ids,

801 - input_ids,

802 - input_batch.cu_num_logits,

803 - self.num_speculative_steps,

793 + sampler_output = self.rejection_sampler(

794 + logits,

795 + input_batch,

796 + # Draft logits are needed for probabilistic rejection sampling.

797 + self.req_states.draft_logits[input_batch.idx_mapping]

798 + if self.req_states.draft_logits is not None

799 + else None,

800 )

805 - sampler_output.sampled_token_ids = sampled_tokens

801

802 # Get the number of sampled and rejected tokens.

803 # For chunked prefills, num_sampled and num_rejected are both 0.

804 num_sampled, num_rejected = get_num_sampled_and_rejected(

810 - num_sampled,

805 + sampler_output.num_sampled,

806 input_batch.seq_lens,

807 input_batch.cu_num_logits,

808 input_batch.idx_mapping,

@@ -1105,6 +1100,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1100 self.req_states.next_prefill_tokens,

1101 self.sampler.sampling_states.temperature.gpu,

1102 self.sampler.sampling_states.seeds.gpu,

1103 + self.req_states.draft_logits,

1104 num_tokens_across_dp=num_tokens_across_dp,

1105 )

1106 self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens

vllm/v1/worker/gpu/sample/gumbel.py

+19 -6

@@ -55,6 +55,8 @@ def _gumbel_sample_kernel(

55 local_argmax_stride,

56 local_max_ptr,

57 local_max_stride,

58 + processed_logits_ptr,

59 + processed_logits_stride,

60 logits_ptr,

61 logits_stride,

62 expanded_idx_mapping_ptr,

@@ -79,6 +81,20 @@ def _gumbel_sample_kernel(

81 logits = logits.to(tl.float32)

83 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)

84 + if (temp != 0.0) and APPLY_TEMPERATURE:

85 + # Apply temperature.

86 + # NOTE(woosuk): Match the behavior of _temperature_kernel.

87 + # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.

88 + logits = logits / temp

89 +

90 + # Store the temperature-applied logits.

91 + if processed_logits_ptr is not None:

92 + tl.store(

93 + processed_logits_ptr + req_state_idx * processed_logits_stride + block,

94 + logits,

95 + mask=mask,

96 + )

97 +

98 if temp != 0.0:

99 # Calculate the seed for gumbel noise.

100 seed = tl.load(seeds_ptr + req_state_idx)

@@ -90,12 +106,6 @@ def _gumbel_sample_kernel(

106 u = tl.maximum(u, 1e-7)

107 gumbel_noise = -tl.log(-tl.log(u))

108

93 - # Apply temperature.

94 - if APPLY_TEMPERATURE:

95 - # NOTE(woosuk): Match the behavior of _temperature_kernel.

96 - # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.

97 - logits = logits / temp

98 -

109 # Apply gumbel noise.

110 logits = tl.where(mask, logits + gumbel_noise, float("-inf"))

111

@@ -112,6 +122,7 @@ def gumbel_sample(

122 seed: torch.Tensor, # [max_num_reqs]

123 pos: torch.Tensor, # [num_tokens]

124 apply_temperature: bool,

125 + processed_logits_out: torch.Tensor | None = None, # [num_reqs, vocab_size]

126 ) -> torch.Tensor:

127 num_tokens, vocab_size = logits.shape

128 BLOCK_SIZE = 1024

@@ -133,6 +144,8 @@ def gumbel_sample(

144 local_argmax.stride(0),

145 local_max,

146 local_max.stride(0),

147 + processed_logits_out,

148 + processed_logits_out.stride(0) if processed_logits_out is not None else 0,

149 logits,

150 logits.stride(0),

151 expanded_idx_mapping,

vllm/v1/worker/gpu/sample/output.py

@@ -12,3 +12,4 @@ class SamplerOutput:

12 sampled_token_ids: torch.Tensor

13 logprobs_tensors: LogprobsTensors | None

14 num_nans: torch.Tensor | None

15 + num_sampled: torch.Tensor | None

vllm/v1/worker/gpu/sample/sampler.py

+32 -10

@@ -7,6 +7,7 @@ import torch

7 import vllm.envs as envs

8 from vllm.config.model import LogprobsMode

9 from vllm.sampling_params import SamplingParams

10 +from vllm.v1.worker.gpu.input_batch import InputBatch

11 from vllm.v1.worker.gpu.metrics.logits import get_num_nans

12 from vllm.v1.worker.gpu.sample.bad_words import BadWordsState

13 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

@@ -56,13 +57,15 @@ class Sampler:

57 def __call__(

58 self,

59 logits: torch.Tensor,

59 - expanded_idx_mapping: torch.Tensor,

60 - idx_mapping_np: np.ndarray,

61 - cu_num_logits_np: np.ndarray,

62 - pos: torch.Tensor,

63 - input_ids: torch.Tensor,

64 - expanded_local_pos: torch.Tensor,

60 + input_batch: InputBatch,

61 ) -> SamplerOutput:

62 + expanded_idx_mapping = input_batch.expanded_idx_mapping

63 + idx_mapping_np = input_batch.idx_mapping_np

64 + cu_num_logits_np = input_batch.cu_num_logits_np

65 + expanded_local_pos = input_batch.expanded_local_pos

66 + pos = input_batch.positions[input_batch.logits_indices]

67 + input_ids = input_batch.input_ids[input_batch.logits_indices]

68 +

69 # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear

70 # that num_nans is computed before applying penalties and temperature.

71 num_nans = get_num_nans(logits) if self.compute_nans else None

@@ -95,10 +98,11 @@ class Sampler:

98 sampled_token_ids=sampled.view(-1, 1),

99 logprobs_tensors=logprobs_tensors,

100 num_nans=num_nans,

101 + num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs),

102 )

103 return sampler_output

104

101 - def sample(

105 + def apply_sampling_params(

106 self,

107 logits: torch.Tensor,

108 expanded_idx_mapping: torch.Tensor,

@@ -106,7 +110,7 @@ class Sampler:

110 pos: torch.Tensor,

111 input_ids: torch.Tensor,

112 expanded_local_pos: torch.Tensor,

109 - ) -> tuple[torch.Tensor, torch.Tensor]:

113 + ) -> torch.Tensor:

114 # Copy logits to a new FP32 tensor.

115 logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)

116

@@ -143,13 +147,31 @@ class Sampler:

147 self.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np)

148

149 # Apply top_k and/or top_p. This might or might not return a new tensor.

146 - logits = self.sampling_states.apply_top_k_top_p(

150 + return self.sampling_states.apply_top_k_top_p(

151 logits, expanded_idx_mapping, idx_mapping_np

152 )

153

154 + def sample(

155 + self,

156 + logits: torch.Tensor,

157 + expanded_idx_mapping: torch.Tensor,

158 + idx_mapping_np: np.ndarray,

159 + pos: torch.Tensor,

160 + input_ids: torch.Tensor,

161 + expanded_local_pos: torch.Tensor,

162 + ) -> tuple[torch.Tensor, torch.Tensor]:

163 + processed_logits = self.apply_sampling_params(

164 + logits,

165 + expanded_idx_mapping,

166 + idx_mapping_np,

167 + pos,

168 + input_ids,

169 + expanded_local_pos,

170 + )

171 +

172 # Sample the next token.

173 sampled = gumbel_sample(

152 - logits,

174 + processed_logits,

175 expanded_idx_mapping,

176 self.sampling_states.temperature.gpu,

177 self.sampling_states.seeds.gpu,

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

+12

@@ -140,6 +140,7 @@ class EagleSpeculator:

140 slot_mappings: dict[str, torch.Tensor] | None,

141 num_tokens_across_dp: torch.Tensor | None,

142 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

143 + draft_logits_out: torch.Tensor | None = None,

144 ) -> None:

145 pos = self.input_buffers.positions[:num_reqs]

146 query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]

@@ -166,6 +167,9 @@ class EagleSpeculator:

167 self.seeds,

168 pos + 1,

169 apply_temperature=True,

170 + processed_logits_out=draft_logits_out[:, step]

171 + if draft_logits_out is not None

172 + else None,

173 )

174 self.draft_tokens[:num_reqs, step] = draft_tokens

175

@@ -219,6 +223,8 @@ class EagleSpeculator:

223 temperature: torch.Tensor,

224 # [max_num_reqs]

225 seeds: torch.Tensor,

226 + # [max_num_reqs, num_speculative_steps, vocab_size]

227 + draft_logits_out: torch.Tensor | None,

228 num_tokens_across_dp: torch.Tensor | None = None,

229 dummy_run: bool = False,

230 skip_attn_for_dummy_run: bool = False,

@@ -271,6 +277,7 @@ class EagleSpeculator:

277 idx_mapping.copy_(input_batch.idx_mapping)

278 self.temperature.copy_(temperature)

279 self.seeds.copy_(seeds)

280 +

281 # Gather the values and copy them to the pre-allocated buffers.

282 pos = self.input_buffers.positions[:num_reqs]

283 torch.gather(input_batch.positions, 0, last_token_indices, out=pos)

@@ -283,7 +290,11 @@ class EagleSpeculator:

290 self.seeds,

291 pos + 1,

292 apply_temperature=True,

293 + processed_logits_out=draft_logits_out[:, 0]

294 + if draft_logits_out is not None

295 + else None,

296 )

297 +

298 if self.num_speculative_steps == 1:

299 # Early exit.

300 return draft_tokens.view(-1, 1)

@@ -365,6 +376,7 @@ class EagleSpeculator:

376 slot_mappings_updated,

377 num_tokens_across_dp=num_tokens_across_dp,

378 cudagraph_runtime_mode=batch_desc.cg_mode,

379 + draft_logits_out=draft_logits_out,

380 )

381 return self.draft_tokens[:num_reqs]

382

vllm/v1/worker/gpu/spec_decode/rejection_sample.py deleted

-62

@@ -1,62 +0,0 @@

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 = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)

52 - num_sampled = cu_num_logits.new_empty(num_reqs)

53 - _rejection_sample_kernel[(num_reqs,)](

54 - sampled,

55 - sampled.stride(0),

56 - num_sampled,

57 - target_sampled,

58 - input_ids,

59 - cu_num_logits,

60 - num_warps=1,

61 - )

62 - return sampled, num_sampled

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py new

+375

@@ -0,0 +1,375 @@

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 +from vllm.v1.worker.gpu.input_batch import InputBatch

7 +from vllm.v1.worker.gpu.metrics.logits import get_num_nans

8 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

9 +from vllm.v1.worker.gpu.sample.output import SamplerOutput

10 +from vllm.v1.worker.gpu.sample.sampler import Sampler

11 +

12 +

13 +@triton.jit

14 +def _strict_rejection_sample_kernel(

15 + sampled_ptr, # [num_reqs, num_speculative_steps + 1]

16 + sampled_stride,

17 + num_sampled_ptr, # [num_reqs]

18 + target_sampled_ptr, # [num_draft_tokens + num_reqs]

19 + input_ids_ptr, # [num_draft_tokens + num_reqs]

20 + cu_num_logits_ptr, # [num_reqs + 1]

21 +):

22 + req_idx = tl.program_id(0)

23 + start_idx = tl.load(cu_num_logits_ptr + req_idx)

24 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)

25 + num_tokens = end_idx - start_idx

26 +

27 + num_sampled = 0

28 + rejected = False

29 + for i in range(num_tokens - 1):

30 + if not rejected:

31 + target_sampled = tl.load(target_sampled_ptr + start_idx + i)

32 + draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1)

33 + tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled)

34 + num_sampled += 1

35 + if target_sampled != draft_sampled:

36 + rejected = True

37 + if not rejected:

38 + target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)

39 + tl.store(

40 + sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled

41 + )

42 + num_sampled += 1

43 + tl.store(num_sampled_ptr + req_idx, num_sampled)

44 +

45 +

46 +def strict_rejection_sample(

47 + # [num_draft_tokens + num_reqs]

48 + target_sampled: torch.Tensor,

49 + # [num_draft_tokens + num_reqs]

50 + draft_sampled: torch.Tensor,

51 + # [num_reqs + 1]

52 + cu_num_logits: torch.Tensor,

53 + num_speculative_steps,

54 +) -> tuple[torch.Tensor, torch.Tensor]:

55 + num_reqs = cu_num_logits.shape[0] - 1

56 + sampled = torch.empty(

57 + num_reqs,

58 + num_speculative_steps + 1,

59 + dtype=target_sampled.dtype,

60 + device=target_sampled.device,

61 + )

62 + num_sampled = torch.empty(

63 + num_reqs,

64 + dtype=torch.int32,

65 + device=target_sampled.device,

66 + )

67 + _strict_rejection_sample_kernel[(num_reqs,)](

68 + sampled,

69 + sampled.stride(0),

70 + num_sampled,

71 + target_sampled,

72 + draft_sampled,

73 + cu_num_logits,

74 + num_warps=1,

75 + )

76 + return sampled, num_sampled

77 +

78 +

79 +@triton.jit

80 +def _probabilistic_rejection_sample_kernel(

81 + # [num_reqs, num_speculative_steps + 1]

82 + sampled_ptr,

83 + sampled_stride,

84 + # [num_reqs]

85 + rejected_steps_ptr,

86 + # [num_logits]

87 + draft_sampled_ptr,

88 + # [num_logits, V]

89 + target_probs_ptr,

90 + target_probs_stride,

91 + # [num_reqs, num_speculative_steps, V]

92 + draft_probs_ptr,

93 + draft_probs_stride_0,

94 + draft_probs_stride_1,

95 + # [num_reqs + 1]

96 + cu_num_logits_ptr,

97 + # [num_logits]

98 + pos_ptr,

99 + # [num_reqs]

100 + idx_mapping_ptr,

101 + # [num_reqs]

102 + seeds_ptr,

103 +):

104 + req_idx = tl.program_id(0)

105 + start_idx = tl.load(cu_num_logits_ptr + req_idx)

106 + num_tokens = tl.load(cu_num_logits_ptr + req_idx + 1) - start_idx

107 + seed = tl.load(seeds_ptr + tl.load(idx_mapping_ptr + req_idx))

108 +

109 + rejected_step = 0

110 + accepted = True

111 + for i in range(num_tokens - 1):

112 + if accepted:

113 + draft_sampled = tl.load(draft_sampled_ptr + start_idx + i + 1)

114 + target_prob = tl.load(

115 + target_probs_ptr + (start_idx + i) * target_probs_stride + draft_sampled

116 + )

117 + draft_prob = tl.load(

118 + draft_probs_ptr

119 + + req_idx * draft_probs_stride_0

120 + + i * draft_probs_stride_1

121 + + draft_sampled

122 + )

123 + pos = tl.load(pos_ptr + start_idx + i)

124 + u = tl.sum(tl.rand(seed, pos + tl.arange(0, 1)))

125 + accepted &= target_prob > u * draft_prob

126 + tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)

127 + rejected_step += accepted

128 + tl.store(rejected_steps_ptr + req_idx, rejected_step)

129 +

130 +

131 +@triton.jit

132 +def _compute_residual_logits_kernel(

133 + # [num_reqs, V]

134 + residual_logits_ptr,

135 + residual_logits_stride,

136 + # [num_reqs]

137 + residual_pos_ptr,

138 + # [num_logits, V]

139 + target_logits_ptr,

140 + target_logits_stride,

141 + # [num_logits, V]

142 + target_probs_ptr,

143 + target_probs_stride,

144 + # [num_reqs, num_speculative_steps, V]

145 + draft_probs_ptr,

146 + draft_probs_stride_0,

147 + draft_probs_stride_1,

148 + # [num_reqs]

149 + rejected_step_ptr,

150 + # [num_reqs + 1]

151 + cu_num_logits_ptr,

152 + # [num_logits]

153 + pos_ptr,

154 + vocab_size,

155 + BLOCK_SIZE: tl.constexpr,

156 +):

157 + req_idx = tl.program_id(0)

158 + block_idx = tl.program_id(1)

159 +

160 + start_idx = tl.load(cu_num_logits_ptr + req_idx)

161 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)

162 + rejected_draft_step = tl.load(rejected_step_ptr + req_idx)

163 + rejected_logit_idx = start_idx + rejected_draft_step

164 +

165 + block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

166 + mask = block_offsets < vocab_size

167 +

168 + if rejected_logit_idx < end_idx - 1:

169 + target_probs = tl.load(

170 + target_probs_ptr + rejected_logit_idx * target_probs_stride + block_offsets,

171 + mask=mask,

172 + other=0.0,

173 + )

174 + draft_probs = tl.load(

175 + draft_probs_ptr

176 + + req_idx * draft_probs_stride_0

177 + + rejected_draft_step * draft_probs_stride_1

178 + + block_offsets,

179 + mask=mask,

180 + other=0.0,

181 + )

182 + residual_probs = tl.maximum(target_probs - draft_probs, 0.0)

183 + residual_logits = tl.log(residual_probs)

184 + else:

185 + # This is a bonus token. Directly return the target logits.

186 + residual_logits = tl.load(

187 + target_logits_ptr

188 + + rejected_logit_idx * target_logits_stride

189 + + block_offsets,

190 + mask=mask,

191 + other=0.0,

192 + )

193 +

194 + tl.store(

195 + residual_logits_ptr + req_idx * residual_logits_stride + block_offsets,

196 + residual_logits,

197 + mask=mask,

198 + )

199 +

200 + # First block computes the residual logit positions.

201 + if block_idx == 0:

202 + pos_val = tl.load(pos_ptr + rejected_logit_idx)

203 + tl.store(residual_pos_ptr + req_idx, pos_val)

204 +

205 +

206 +def probabilistic_rejection_sample(

207 + # [num_draft_tokens + num_reqs, V]

208 + target_logits: torch.Tensor,

209 + # [num_reqs, num_speculative_steps, V]

210 + draft_logits: torch.Tensor,

211 + # [num_draft_tokens + num_reqs]

212 + draft_sampled: torch.Tensor,

213 + # [num_reqs + 1]

214 + cu_num_logits: torch.Tensor,

215 + # [num_logits]

216 + pos: torch.Tensor,

217 + # [num_reqs]

218 + idx_mapping: torch.Tensor,

219 + temperature,

220 + seeds,

221 + num_speculative_steps,

222 +) -> tuple[torch.Tensor, torch.Tensor]:

223 + num_reqs = cu_num_logits.shape[0] - 1

224 + device = target_logits.device

225 + vocab_size = target_logits.shape[-1]

226 +

227 + # Compute target and draft probs.

228 + target_probs = torch.softmax(target_logits, dim=-1)

229 + draft_probs = torch.softmax(draft_logits, dim=-1)

230 +

231 + # Rejection sample.

232 + # [num_reqs, num_speculative_steps + 1]

233 + sampled = torch.empty(

234 + num_reqs,

235 + num_speculative_steps + 1,

236 + dtype=torch.int64,

237 + device=device,

238 + )

239 + # [num_reqs]

240 + rejected_steps = torch.empty(

241 + num_reqs,

242 + dtype=torch.int64,

243 + device=device,

244 + )

245 + _probabilistic_rejection_sample_kernel[(num_reqs,)](

246 + sampled,

247 + sampled.stride(0),

248 + rejected_steps,

249 + draft_sampled,

250 + target_probs,

251 + target_probs.stride(0),

252 + draft_probs,

253 + draft_probs.stride(0),

254 + draft_probs.stride(1),

255 + cu_num_logits,

256 + pos,

257 + idx_mapping,

258 + seeds,

259 + num_warps=1,

260 + )

261 +

262 + # Compute the logits and positions to resample the rejected/bonus

263 + # tokens from.

264 + # [num_reqs, vocab_size]

265 + residual_logits = torch.empty(

266 + num_reqs,

267 + vocab_size,

268 + dtype=target_logits.dtype,

269 + device=device,

270 + )

271 + # [num_reqs]

272 + residual_pos = torch.empty(

273 + num_reqs,

274 + dtype=pos.dtype,

275 + device=device,

276 + )

277 + BLOCK_SIZE = 1024

278 + num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

279 + _compute_residual_logits_kernel[(num_reqs, num_blocks)](

280 + residual_logits,

281 + residual_logits.stride(0),

282 + residual_pos,

283 + target_logits,

284 + target_logits.stride(0),

285 + target_probs,

286 + target_probs.stride(0),

287 + draft_probs,

288 + draft_probs.stride(0),

289 + draft_probs.stride(1),

290 + rejected_steps,

291 + cu_num_logits,

292 + pos,

293 + vocab_size,

294 + BLOCK_SIZE=BLOCK_SIZE,

295 + )

296 +

297 + # Gumbel sample tokens from the residual distribution.

298 + resampled = gumbel_sample(

299 + residual_logits,

300 + idx_mapping,

301 + temperature,

302 + seeds,

303 + residual_pos,

304 + apply_temperature=False,

305 + )

306 + sampled.scatter_(1, rejected_steps.unsqueeze(1), resampled.unsqueeze(1))

307 +

308 + return sampled, rejected_steps + 1

309 +

310 +

311 +class RejectionSampler:

312 + def __init__(

313 + self,

314 + sampler: Sampler,

315 + num_speculative_steps,

316 + use_strict_rejection_sampling: bool = True,

317 + ):

318 + self.sampler = sampler

319 + self.num_speculative_steps = num_speculative_steps

320 + self.use_strict_rejection_sampling = use_strict_rejection_sampling

321 +

322 + def __call__(

323 + self,

324 + logits: torch.Tensor,

325 + input_batch: InputBatch,

326 + draft_logits: torch.Tensor | None = None,

327 + ) -> SamplerOutput:

328 + draft_sampled = input_batch.input_ids[input_batch.logits_indices]

329 + # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear

330 + # that num_nans is computed before applying penalties and temperature.

331 + num_nans = get_num_nans(logits) if self.sampler.compute_nans else None

332 +

333 + if self.use_strict_rejection_sampling:

334 + sampler_output = self.sampler(

335 + logits,

336 + input_batch,

337 + )

338 + logprobs_tensors = sampler_output.logprobs_tensors

339 + sampled, num_sampled = strict_rejection_sample(

340 + sampler_output.sampled_token_ids.view(-1),

341 + draft_sampled,

342 + input_batch.cu_num_logits,

343 + self.num_speculative_steps,

344 + )

345 + else:

346 + assert draft_logits is not None

347 + pos = input_batch.positions[input_batch.logits_indices]

348 + processed_logits = self.sampler.apply_sampling_params(

349 + logits,

350 + input_batch.expanded_idx_mapping,

351 + input_batch.idx_mapping_np,

352 + pos,

353 + draft_sampled,

354 + input_batch.expanded_local_pos,

355 + )

356 + # TODO (TheEpicDolphin): Return logprobs for sampled token ids.

357 + logprobs_tensors = None

358 + sampled, num_sampled = probabilistic_rejection_sample(

359 + processed_logits,

360 + draft_logits,

361 + draft_sampled,

362 + input_batch.cu_num_logits,

363 + pos,

364 + input_batch.idx_mapping,

365 + self.sampler.sampling_states.temperature.gpu,

366 + self.sampler.sampling_states.seeds.gpu,

367 + self.num_speculative_steps,

368 + )

369 +

370 + return SamplerOutput(

371 + sampled_token_ids=sampled,

372 + logprobs_tensors=logprobs_tensors,

373 + num_nans=num_nans,

374 + num_sampled=num_sampled,

375 + )

vllm/v1/worker/gpu/states.py

+15

@@ -15,6 +15,8 @@ class RequestState:

15 num_speculative_steps: int,

16 vocab_size: int,

17 device: torch.device,

18 + model_dtype: torch.dtype,

19 + cache_draft_logits: bool,

20 ):

21 self.max_num_reqs = max_num_reqs

22 self.max_model_len = max_model_len

@@ -70,6 +72,19 @@ class RequestState:

72 dtype=torch.int64,

73 device=device,

74 )

75 + # Draft token logits.

76 + # NOTE: This tensor maintains the "processed" logits after applying temperature,

77 + # top-p, etc.

78 + self.draft_logits: torch.Tensor | None = None

79 + if cache_draft_logits:

80 + self.draft_logits = torch.zeros(

81 + self.max_num_reqs,

82 + self.num_speculative_steps,

83 + self.vocab_size,

84 + dtype=model_dtype,

85 + device=device,

86 + )

87 +

88 self.next_prefill_tokens = torch.zeros(

89 self.max_num_reqs, dtype=torch.int32, device=device

90 )