xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-15 13:05:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

[V1][Minor] Simplify rejection sampler's parse_output (#15741)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2025-03-29 09:25:17 -07:00 committed by GitHub
parent c67abd614f
commit 2bc4be4e32
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 3 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/v1/sample/rejection_sampler.py
View File

@ -107,7 +107,6 @@ class RejectionSampler(nn.Module):
@staticmethod
def parse_output(
output_token_ids: torch.Tensor,
ignored_req_idxs: list[int],
vocab_size: int,
) -> list[list[int]]:
"""Parse the output of the rejection sampler.
@ -117,9 +116,6 @@ class RejectionSampler(nn.Module):
[batch_size, max_spec_len + 1]. The rejected tokens are
replaced with `PLACEHOLDER_TOKEN_ID` by the rejection sampler
and will be filtered out in this function.
ignored_req_idxs: The indices of the requests that should not be
sampled. This is usually because the request is still in the
prefill phase.
vocab_size: The size of the vocabulary.
Returns:
@ -129,11 +125,8 @@ class RejectionSampler(nn.Module):
# Create mask for valid tokens.
valid_mask = ((output_token_ids_np != PLACEHOLDER_TOKEN_ID) &
(output_token_ids_np < vocab_size))
ignored_req_idx_set = set(ignored_req_idxs)
outputs = [
row[valid_mask[i]].tolist()
if i not in ignored_req_idx_set else []
for i, row in enumerate(output_token_ids_np)
]
return outputs

7
vllm/v1/worker/gpu_model_runner.py
View File

@ -1121,16 +1121,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
if max_gen_len == 1:
# No spec decode tokens.
valid_sampled_token_ids = sampled_token_ids.tolist()
# Mask out the sampled tokens that should not be sampled.
for i in discard_sampled_tokens_req_indices:
valid_sampled_token_ids[i].clear()
else:
# Includes spec decode tokens.
valid_sampled_token_ids = self.rejection_sampler.parse_output(
sampled_token_ids,
discard_sampled_tokens_req_indices,
self.input_batch.vocab_size,
)
# Mask out the sampled tokens that should not be sampled.
for i in discard_sampled_tokens_req_indices:
valid_sampled_token_ids[i].clear()
if not self.use_spec_decode:
spec_token_ids = None