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 15, 13:30 UTC
[Core] Speed up decode by remove synchronizing operation in sampler (#16436)
Signed-off-by: Chanh Nguyen <cnguyen@linkedin.com> Co-authored-by: Chanh Nguyen <cnguyen@linkedin.com>
Chanh Nguyen committed
Apr 21, 2025 at 11:18 UTC
299ebb62b269ce167eb1c71b5e39a1dc1f65ce1c
1 file changed
+9
-4
vllm/model_executor/layers/utils.py
+9
-4
@@ -47,10 +47,15 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
47
output_tokens_tensor, vocab_size, num_seqs)
48
repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat(
49
1, vocab_size)
50
- logits[logits > 0] /= torch.where(prompt_mask | output_mask,
51
- repetition_penalties, 1.0)[logits > 0]
52
- logits[logits <= 0] *= torch.where(prompt_mask | output_mask,
53
- repetition_penalties, 1.0)[logits <= 0]
50
+
51
+ # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
52
+ penalties = torch.where(prompt_mask | output_mask, repetition_penalties,
53
+ 1.0)
54
+
55
+ # If logits are positive, divide by penalty, otherwise multiply by penalty.
56
+ scaling = torch.where(logits > 0, 1.0 / penalties, penalties)
57
+ logits *= scaling
58
+
59
# We follow the definition in OpenAI API.
60
# Refer to https://platform.openai.com/docs/api-reference/parameter-details
61
logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts