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, 08:30 UTC
[Model Runner V2] Fuse penalties and temperature into single kernel (#29720)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 29, 2025 at 02:29 UTC
f223ed4181f003301a115950178746f42b668903
3 files changed
+59
-45
vllm/v1/worker/gpu/sample/gumbel.py
+3
-2
@@ -45,8 +45,9 @@ def _gumbel_sample_kernel(
45
46
# Apply temperature.
47
if APPLY_TEMPERATURE:
48
- # NOTE(woosuk): Use div_rn to match the behavior of torch.
49
- logits = tl.div_rn(logits, temp)
48
+ # NOTE(woosuk): Match the behavior of _penalties_and_temperature_kernel.
49
+ # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.
50
+ logits = logits / temp
51
52
# Apply gumbel noise.
53
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
vllm/v1/worker/gpu/sample/penalties.py
+42
-25
@@ -7,12 +7,13 @@ from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
7
8
9
@triton.jit
10
-def _penalties_kernel(
10
+def _penalties_and_temperature_kernel(
11
logits_ptr,
12
logits_stride,
13
repetition_penalty_ptr,
14
frequency_penalty_ptr,
15
presence_penalty_ptr,
16
+ temperature_ptr,
17
idx_mapping_ptr,
18
prompt_bin_counts_ptr,
19
prompt_bin_counts_stride,
@@ -25,12 +26,16 @@ def _penalties_kernel(
26
rep_penalty = tl.load(repetition_penalty_ptr + batch_idx)
27
freq_penalty = tl.load(frequency_penalty_ptr + batch_idx)
28
pres_penalty = tl.load(presence_penalty_ptr + batch_idx)
29
+ temperature = tl.load(temperature_ptr + batch_idx)
30
+ temperature = tl.where(temperature == 0.0, 1.0, temperature)
31
32
use_rep_penalty = rep_penalty != 1.0
33
use_freq_penalty = freq_penalty != 0.0
34
use_pres_penalty = pres_penalty != 0.0
32
- if not (use_rep_penalty or use_freq_penalty or use_pres_penalty):
33
- # No penalties to apply. Early return.
35
+ use_penalty = use_rep_penalty or use_freq_penalty or use_pres_penalty
36
+ use_temperature = temperature != 1.0
37
+ if not (use_penalty or use_temperature):
38
+ # Early return to avoid loading logits.
39
return
40
41
block_idx = tl.program_id(1)
@@ -39,42 +44,54 @@ def _penalties_kernel(
44
logits = tl.load(logits_ptr + batch_idx * logits_stride + block, mask=mask)
45
logits = logits.to(tl.float32)
46
42
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
43
- output_bin_counts = tl.load(
44
- output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
45
- mask=mask,
46
- )
47
-
48
- # Apply repetition penalties.
49
- if use_rep_penalty:
50
- prompt_bin_counts = tl.load(
51
- prompt_bin_counts_ptr + req_state_idx * prompt_bin_counts_stride + block,
47
+ if use_penalty:
48
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
49
+ output_bin_counts = tl.load(
50
+ output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
51
mask=mask,
52
)
54
- # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
55
- scale = tl.where((prompt_bin_counts + output_bin_counts) > 0, rep_penalty, 1.0)
56
- # If logits are positive, divide by penalty, otherwise multiply by penalty.
57
- scale = tl.where(logits > 0, 1.0 / scale, scale)
58
- logits *= scale
59
-
60
- # Apply frequency penalties.
61
- logits -= freq_penalty * output_bin_counts
62
- # Apply presence penalties.
63
- logits -= pres_penalty * (output_bin_counts > 0)
53
+ output_bin_mask = output_bin_counts > 0
54
+
55
+ # Apply repetition penalties.
56
+ if use_rep_penalty:
57
+ prompt_bin_counts = tl.load(
58
+ prompt_bin_counts_ptr
59
+ + req_state_idx * prompt_bin_counts_stride
60
+ + block,
61
+ mask=mask,
62
+ )
63
+ prompt_bin_mask = prompt_bin_counts > 0
64
+ # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
65
+ scale = tl.where(prompt_bin_mask | output_bin_mask, rep_penalty, 1.0)
66
+ # If logits are positive, divide by penalty, otherwise multiply by penalty.
67
+ logits *= tl.where(logits > 0, 1.0 / scale, scale)
68
+
69
+ # Apply frequency penalties.
70
+ logits -= freq_penalty * output_bin_counts
71
+ # Apply presence penalties.
72
+ logits -= pres_penalty * output_bin_mask
73
+
74
+ # Apply temperature.
75
+ logits = logits / temperature
76
+
77
# Store back to logits.
78
tl.store(logits_ptr + batch_idx * logits_stride + block, logits, mask=mask)
79
80
68
-def apply_penalties(logits: torch.Tensor, sampling_metadata: SamplingMetadata) -> None:
81
+def apply_penalties_and_temperature(
82
+ logits: torch.Tensor,
83
+ sampling_metadata: SamplingMetadata,
84
+) -> None:
85
num_reqs, vocab_size = logits.shape
86
BLOCK_SIZE = 8192
87
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
72
- _penalties_kernel[(num_reqs, num_blocks)](
88
+ _penalties_and_temperature_kernel[(num_reqs, num_blocks)](
89
logits,
90
logits.stride(0),
91
sampling_metadata.repetition_penalty,
92
sampling_metadata.frequency_penalty,
93
sampling_metadata.presence_penalty,
94
+ sampling_metadata.temperature,
95
sampling_metadata.idx_mapping,
96
sampling_metadata.prompt_bin_counts,
97
sampling_metadata.prompt_bin_counts.stride(0),
vllm/v1/worker/gpu/sample/sampler.py
+14
-18
@@ -9,7 +9,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
9
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
10
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
11
from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
12
-from vllm.v1.worker.gpu.sample.penalties import apply_penalties
12
+from vllm.v1.worker.gpu.sample.penalties import apply_penalties_and_temperature
13
14
15
class Sampler:
@@ -26,22 +26,19 @@ class Sampler:
26
logits: torch.Tensor,
27
sampling_metadata: SamplingMetadata,
28
) -> SamplerOutput:
29
+ sampled, processed_logits = self.sample(logits, sampling_metadata)
30
if sampling_metadata.max_num_logprobs is not None:
30
- if self.logprobs_mode == "processed_logprobs":
31
- sampled, logits = self.sample(
32
- logits, sampling_metadata, return_logits=True
33
- )
34
- else:
35
- assert self.logprobs_mode == "raw_logprobs"
36
- sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
37
-
31
+ logits = (
32
+ processed_logits
33
+ if self.logprobs_mode == "processed_logprobs"
34
+ else logits
35
+ )
36
logprobs_tensors = compute_topk_logprobs(
37
logits,
38
sampling_metadata.max_num_logprobs,
39
sampled,
40
)
41
else:
44
- sampled, _ = self.sample(logits, sampling_metadata, return_logits=False)
42
logprobs_tensors = None
43
44
# These are GPU tensors.
@@ -58,16 +55,15 @@ class Sampler:
55
self,
56
logits: torch.Tensor,
57
sampling_metadata: SamplingMetadata,
61
- return_logits: bool = False,
62
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
63
- is_greedy = sampling_metadata.temperature == 0
64
- temp = torch.where(is_greedy, 1.0, sampling_metadata.temperature)
65
- logits = logits / temp.view(-1, 1)
58
+ ) -> tuple[torch.Tensor, torch.Tensor]:
59
+ # Copy logits to a new FP32 tensor.
60
+ logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
61
+
62
+ # Apply penalties and temperature in place.
63
+ apply_penalties_and_temperature(logits, sampling_metadata)
64
logits = apply_top_k_top_p(
65
logits, sampling_metadata.top_k, sampling_metadata.top_p
66
)
69
- # Apply penalties in place.
70
- apply_penalties(logits, sampling_metadata)
67
68
sampled = gumbel_sample(
69
logits,
@@ -76,4 +72,4 @@ class Sampler:
72
sampling_metadata.pos,
73
apply_temperature=False,
74
)
79
- return sampled, logits if return_logits else None
75
+ return sampled, logits