xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-22 13:24:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
LiuXiaoxuanPKU 2025-04-01 15:46:41 -07:00
parent 8fcd4d18e0
commit b484b79504
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -195,7 +195,7 @@ def rejection_sample(
num_warps=1, num_warps=1,
) )
if sampling_metadata.all_greedy: if sampling_metadata.all_greedy:
return output_token_ids return output_token_ids, output_probs
# Generate uniform probabilities for rejection sampling. # Generate uniform probabilities for rejection sampling.
# [num_tokens] # [num_tokens]
@ -475,8 +475,8 @@ def rejection_greedy_sample_kernel(
if draft_token_id != target_argmax_id: if draft_token_id != target_argmax_id:
# Reject. # Reject.
rejected = True rejected = True
tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, tl.store(output_probs_ptr + req_idx * (max_spec_len + 1) + pos,
not rejected) draft_token_id == target_argmax_id)
if not rejected: if not rejected:
# If all tokens are accepted, append the bonus token. # If all tokens are accepted, append the bonus token.