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, 19:15 UTC
[Model Runner V2] Some code simplification (#36929)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 12, 2026 at 17:59 UTC
cd32d6f5868a040430a88c4423e3307116feb433
3 files changed
+17
-59
vllm/config/speculative.py
+1
-4
@@ -57,10 +57,7 @@ SpeculativeMethod = Literal[
57
EagleModelTypes,
58
NgramGPUTypes,
59
]
60
-RejectionSampleMethod = Literal[
61
- "strict",
62
- "probabilistic",
63
-]
60
+RejectionSampleMethod = Literal["strict", "probabilistic"]
61
62
63
@config
vllm/v1/worker/gpu/sample/gumbel.py
+3
-13
@@ -81,7 +81,7 @@ def _gumbel_sample_kernel(
81
logits = logits.to(tl.float32)
82
83
temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
84
- if (temp != 0.0) and APPLY_TEMPERATURE:
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.
@@ -127,18 +127,8 @@ def gumbel_sample(
127
num_tokens, vocab_size = logits.shape
128
BLOCK_SIZE = 1024
129
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
130
- local_argmax = torch.empty(
131
- num_tokens,
132
- num_blocks,
133
- dtype=torch.int64,
134
- device=logits.device,
135
- )
136
- local_max = torch.empty(
137
- num_tokens,
138
- num_blocks,
139
- dtype=torch.float32,
140
- device=logits.device,
141
- )
130
+ local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
131
+ local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float32)
132
_gumbel_sample_kernel[(num_tokens, num_blocks)](
133
local_argmax,
134
local_argmax.stride(0),
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+13
-42
@@ -53,17 +53,8 @@ def strict_rejection_sample(
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
- )
56
+ sampled = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)
57
+ num_sampled = target_sampled.new_empty(num_reqs, dtype=torch.int32)
58
_strict_rejection_sample_kernel[(num_reqs,)](
59
sampled,
60
sampled.stride(0),
@@ -216,12 +207,11 @@ def probabilistic_rejection_sample(
207
pos: torch.Tensor,
208
# [num_reqs]
209
idx_mapping: torch.Tensor,
219
- temperature,
220
- seeds,
221
- num_speculative_steps,
210
+ temperature: torch.Tensor,
211
+ seed: torch.Tensor,
212
+ num_speculative_steps: int,
213
) -> tuple[torch.Tensor, torch.Tensor]:
214
num_reqs = cu_num_logits.shape[0] - 1
224
- device = target_logits.device
215
vocab_size = target_logits.shape[-1]
216
217
# Compute target and draft probs.
@@ -230,18 +220,11 @@ def probabilistic_rejection_sample(
220
221
# Rejection sample.
222
# [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,
223
+ sampled = draft_sampled.new_empty(
224
+ num_reqs, num_speculative_steps + 1, dtype=torch.int64
225
)
226
# [num_reqs]
240
- rejected_steps = torch.empty(
241
- num_reqs,
242
- dtype=torch.int64,
243
- device=device,
244
- )
227
+ rejected_steps = sampled.new_empty(num_reqs)
228
_probabilistic_rejection_sample_kernel[(num_reqs,)](
229
sampled,
230
sampled.stride(0),
@@ -255,25 +238,16 @@ def probabilistic_rejection_sample(
238
cu_num_logits,
239
pos,
240
idx_mapping,
258
- seeds,
241
+ seed,
242
num_warps=1,
243
)
244
245
# Compute the logits and positions to resample the rejected/bonus
246
# tokens from.
247
# [num_reqs, vocab_size]
265
- residual_logits = torch.empty(
266
- num_reqs,
267
- vocab_size,
268
- dtype=target_logits.dtype,
269
- device=device,
270
- )
248
+ residual_logits = target_logits.new_empty(num_reqs, vocab_size)
249
# [num_reqs]
272
- residual_pos = torch.empty(
273
- num_reqs,
274
- dtype=pos.dtype,
275
- device=device,
276
- )
250
+ residual_pos = pos.new_empty(num_reqs)
251
BLOCK_SIZE = 1024
252
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
253
_compute_residual_logits_kernel[(num_reqs, num_blocks)](
@@ -299,7 +273,7 @@ def probabilistic_rejection_sample(
273
residual_logits,
274
idx_mapping,
275
temperature,
302
- seeds,
276
+ seed,
277
residual_pos,
278
apply_temperature=False,
279
)
@@ -331,10 +305,7 @@ class RejectionSampler:
305
num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
306
307
if self.use_strict_rejection_sampling:
334
- sampler_output = self.sampler(
335
- logits,
336
- input_batch,
337
- )
308
+ sampler_output = self.sampler(logits, input_batch)
309
logprobs_tensors = sampler_output.logprobs_tensors
310
sampled, num_sampled = strict_rejection_sample(
311
sampler_output.sampled_token_ids.view(-1),