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 3, 21:30 UTC
[Bugfix][Performance Improvement] Improve penalties triton kernel performance (#40657)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
May 11, 2026 at 22:47 UTC
4e498b5e5c07480cfb8c046128f0ef8d9a60d8ef
2 files changed
+8
-16
vllm/v1/worker/gpu/sample/penalties.py
+8
-15
@@ -85,7 +85,6 @@ class PenaltiesState:
85
idx_mapping_np: np.ndarray,
86
input_ids: torch.Tensor,
87
expanded_local_pos: torch.Tensor,
88
- num_speculative_tokens: int,
88
) -> None:
89
if not np.any(self.use_penalty[idx_mapping_np]):
90
# No request uses penalties. Skip the kernel launch.
@@ -101,7 +100,6 @@ class PenaltiesState:
100
self.presence_penalty.gpu,
101
self.prompt_bin_mask,
102
self.output_bin_counts,
104
- num_speculative_tokens,
103
)
104
105
@@ -121,7 +119,6 @@ def _penalties_kernel(
119
output_bin_counts_stride,
120
vocab_size,
121
BLOCK_SIZE: tl.constexpr,
124
- MAX_SPEC_LEN: tl.constexpr,
122
):
123
token_idx = tl.program_id(0)
124
req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
@@ -149,18 +146,16 @@ def _penalties_kernel(
146
other=0,
147
)
148
152
- # Compute cumulative draft_counts from previous positions in this request
149
+ # Accumulate draft token counts from previous positions directly into
150
+ # output_bin_counts (preserves its native tensor layout, avoiding an
151
+ # expensive shared-memory layout conversion after the loop).
152
pos = tl.load(expanded_local_pos_ptr + token_idx)
153
start_idx = token_idx - pos
155
- draft_counts = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
156
- for prev_pos in tl.static_range(MAX_SPEC_LEN):
157
- if prev_pos < pos:
158
- prev_token = tl.load(token_ids_ptr + start_idx + prev_pos + 1)
159
- token_match = block == prev_token
160
- draft_counts = draft_counts + token_match.to(tl.int32)
161
-
162
- # Total counts = base output counts + cumulative draft counts
163
- output_bin_counts = base_output_counts + draft_counts
154
+ output_bin_counts = base_output_counts
155
+ for prev_pos in tl.range(pos):
156
+ prev_token = tl.load(token_ids_ptr + start_idx + prev_pos + 1)
157
+ token_match = block == prev_token
158
+ output_bin_counts = output_bin_counts + token_match.to(tl.int32)
159
output_bin_mask = output_bin_counts > 0
160
161
# Apply repetition penalties.
@@ -198,7 +193,6 @@ def apply_penalties(
193
presence_penalty: torch.Tensor,
194
prompt_bin_mask: torch.Tensor,
195
output_bin_counts: torch.Tensor,
201
- num_speculative_tokens: int,
196
) -> None:
197
num_tokens, vocab_size = logits.shape
198
BLOCK_SIZE = 8192
@@ -218,7 +212,6 @@ def apply_penalties(
212
output_bin_counts.stride(0),
213
vocab_size,
214
BLOCK_SIZE=BLOCK_SIZE,
221
- MAX_SPEC_LEN=num_speculative_tokens,
215
)
216
217
vllm/v1/worker/gpu/sample/sampler.py
-1
@@ -142,7 +142,6 @@ class Sampler:
142
idx_mapping_np,
143
input_ids,
144
expanded_local_pos,
145
- self.num_speculative_tokens,
145
)
146
147
# Apply bad words masking in place.