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, 17:15 UTC
[V1][Minor] Simplify rejection sampler's parse_output (#15741)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 29, 2025 at 09:25 UTC
2bc4be4e32a42a439f7aad3752b96a20e7c34938
2 files changed
+3
-11
vllm/v1/sample/rejection_sampler.py
-7
@@ -107,7 +107,6 @@ class RejectionSampler(nn.Module):
107
@staticmethod
108
def parse_output(
109
output_token_ids: torch.Tensor,
110
- ignored_req_idxs: list[int],
110
vocab_size: int,
111
) -> list[list[int]]:
112
"""Parse the output of the rejection sampler.
@@ -117,9 +116,6 @@ class RejectionSampler(nn.Module):
116
[batch_size, max_spec_len + 1]. The rejected tokens are
117
replaced with `PLACEHOLDER_TOKEN_ID` by the rejection sampler
118
and will be filtered out in this function.
120
- ignored_req_idxs: The indices of the requests that should not be
121
- sampled. This is usually because the request is still in the
122
- prefill phase.
119
vocab_size: The size of the vocabulary.
120
121
Returns:
@@ -129,11 +125,8 @@ class RejectionSampler(nn.Module):
125
# Create mask for valid tokens.
126
valid_mask = ((output_token_ids_np != PLACEHOLDER_TOKEN_ID) &
127
(output_token_ids_np < vocab_size))
132
-
133
- ignored_req_idx_set = set(ignored_req_idxs)
128
outputs = [
129
row[valid_mask[i]].tolist()
136
- if i not in ignored_req_idx_set else []
130
for i, row in enumerate(output_token_ids_np)
131
]
132
return outputs
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -1121,16 +1121,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1121
if max_gen_len == 1:
1122
# No spec decode tokens.
1123
valid_sampled_token_ids = sampled_token_ids.tolist()
1124
- # Mask out the sampled tokens that should not be sampled.
1125
- for i in discard_sampled_tokens_req_indices:
1126
- valid_sampled_token_ids[i].clear()
1124
else:
1125
# Includes spec decode tokens.
1126
valid_sampled_token_ids = self.rejection_sampler.parse_output(
1127
sampled_token_ids,
1131
- discard_sampled_tokens_req_indices,
1128
self.input_batch.vocab_size,
1129
)
1130
+ # Mask out the sampled tokens that should not be sampled.
1131
+ for i in discard_sampled_tokens_req_indices:
1132
+ valid_sampled_token_ids[i].clear()
1133
1134
if not self.use_spec_decode:
1135
spec_token_ids = None