[V1][spec decode] return logprobs for spec decoding

Signed-off-by: Giancarlo Delfin <gdelfin@meta.com>

Author: TheEpicDolphin

vllm/v1/sample/rejection_sampler.py

@@ -51,9 +51,9 @@ class RejectionSampler(nn.Module):
     def __init__(self, sampler: Sampler):
         super().__init__()
         self.sampler = sampler
-        self.return_processed_logprobs = self.sampler.logprobs_mode.startswith(
-            "processed"
-        )
+        logprobs_mode = self.sampler.logprobs_mode
+        self.is_processed_logprobs_mode = logprobs_mode.startswith("processed")
+        self.is_logits_logprobs_mode = logprobs_mode.endswith("logits")
 
     def forward(
         self,
@@ -107,7 +107,7 @@ def forward(
             # Override the logprobs mode to return logits because they are
             # needed later to compute the accepted token logprobs.
             logprobs_mode_override="processed_logits"
-            if self.return_processed_logprobs
+            if self.is_processed_logprobs_mode
             else "raw_logits",
         )
         bonus_logits = bonus_sampler_output.logprobs_tensors.logprobs
@@ -150,7 +150,7 @@ def forward(
                 sampling_metadata,
                 metadata,
                 logits,
-                target_logits if self.return_processed_logprobs else raw_target_logits,
+                target_logits if self.is_processed_logprobs_mode else raw_target_logits,
                 bonus_logits,
                 output_token_ids,
             ),
@@ -190,7 +190,7 @@ def _get_logprobs_tensors(
         accepted_logits = final_logits[accepted_logit_indices]
         accepted_logprobs = (
             accepted_logits
-            if self.logprobs_mode.endswith("logits")
+            if self.is_logits_logprobs_mode
             else self.sampler.compute_logprobs(accepted_logits)
         )
         accepted_tokens = sampled_token_ids[accepted_mask]