update return type

Signed-off-by: NickLucche <[email protected]>

Author: NickLucche

vllm/v1/worker/tpu_model_runner.py

@@ -787,7 +787,8 @@ def execute_model(
             logits = self.structured_decode(require_struct_decoding,
                                             grammar_bitmask_padded, logits,
                                             arange)
-        selected_token_ids, logprobs = self.sample_from_logits(logits, tpu_sampling_metadata)
+        selected_token_ids, logprobs = self.sample_from_logits(
+            logits, tpu_sampling_metadata)
         # Remove padding on cpu and keep dynamic op outside of xla graph.
         selected_token_ids = selected_token_ids.cpu()[:num_reqs]
         logprobs_lists = logprobs.tolists() \
@@ -1254,7 +1255,8 @@ def compute_logits(self,
     @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
     def sample_from_logits(
             self, logits: torch.Tensor,
-            sampling_metadata: TPUSupportedSamplingMetadata) -> torch.Tensor:
+            sampling_metadata: TPUSupportedSamplingMetadata) -> \
+                tuple[torch.Tensor, Optional[LogprobsTensors]]:
         """
         Sample with xla-friendly function. This function is to be traced 
         separately from `forward` for lighter compilation overhead.