From b484b79504de4d1c35ec737acef6df6529f68359 Mon Sep 17 00:00:00 2001
From: LiuXiaoxuanPKU <lilyliupku@gmail.com>
Date: Tue, 1 Apr 2025 15:46:41 -0700
Subject: [PATCH] fix

---
 vllm/v1/sample/rejection_sampler.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

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