[Bugfix] Fix error where vLLM expects numpy sampled token ids

Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>

Author: kyuyeunk

tpu_inference/runner/tpu_runner.py

@@ -15,7 +15,7 @@
 from flax import nnx
 from jax.experimental import mesh_utils
 from jax.sharding import NamedSharding, PartitionSpec
-from torchax.ops.mappings import j2t_dtype
+from torchax.ops.mappings import j2t, j2t_dtype
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                           has_kv_transfer_group)
@@ -28,7 +28,7 @@
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,
                              DraftTokenIds, KVConnectorOutput, LogprobsLists,
-                             ModelRunnerOutput)
+                             LogprobsTensors, ModelRunnerOutput)
 from vllm.v1.request import Request
 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
 from vllm.v1.worker.kv_connector_model_runner_mixin import \
@@ -122,7 +122,8 @@ def get_output(self) -> ModelRunnerOutput:
             next_tokens_cpu = next_tokens_cpu[self.logits_indices_selector]
         selected_token_ids = np.expand_dims(next_tokens_cpu[:self._num_reqs],
                                             1)
-        valid_sampled_token_ids = selected_token_ids.tolist()
+
+        valid_sampled_token_ids = [token_id for token_id in selected_token_ids]
         for i in self._discard_sampled_tokens_req_indices:
             valid_sampled_token_ids[i].clear()
         self._model_runner_output.sampled_token_ids = valid_sampled_token_ids
@@ -190,7 +191,8 @@ def _substitute_placeholder_token(
     return input_ids.at[token_in_tpu_cur_input_indices].set(update_values)
 
 
-def _reorder_logits_indices(logprobs_lists, logits_indices_selector):
+def _reorder_logits_indices(logprobs_lists: LogprobsLists,
+                            logits_indices_selector: List[int]):
     return LogprobsLists(
         logprob_token_ids=[
             logprobs_lists.logprob_token_ids[i]
@@ -595,7 +597,7 @@ def _modify_prev_results(self):
             next_tokens_cpu = next_tokens_cpu[pre_logits_indices_selector]
         selected_token_ids = np.expand_dims(next_tokens_cpu[:len(pre_req_ids)],
                                             1)
-        valid_sampled_token_ids = selected_token_ids.tolist()
+        valid_sampled_token_ids = [token_id for token_id in selected_token_ids]
 
         # Mask out the sampled tokens that should not be sampled.
         for i in pre_discard_sampled_tokens_req_indices:
@@ -898,7 +900,9 @@ def _sample_from_logits(
             if logits_indices_selector is not None:
                 next_tokens = next_tokens[logits_indices_selector]
             selected_token_ids = np.expand_dims(next_tokens[:num_reqs], 1)
-            valid_sampled_token_ids = selected_token_ids.tolist()
+            valid_sampled_token_ids = [
+                token_id for token_id in selected_token_ids
+            ]
         else:
             valid_sampled_token_ids = self.rejection_sampler.parse_output(
                 next_tokens, self.input_batch.vocab_size,
@@ -975,10 +979,17 @@ def select_local_fn(local_array, local_indices):
         return ret
 
     @staticmethod
-    @functools.partial(jax.jit, static_argnames=("max_logprobs", ))
-    def _compute_and_gather_logprobs(logits, next_tokens, max_logprobs):
-        logprobs = compute_logprobs(logits)
-        return gather_logprobs(logprobs, next_tokens, max_logprobs)
+    def _compute_and_gather_logprobs(logits, next_tokens,
+                                     max_logprobs) -> LogprobsTensors:
+
+        @functools.partial(jax.jit, static_argnames=("max_logprobs", ))
+        def jit_compute_and_gather_logprobs(logits, next_tokens, max_logprobs):
+            logprobs = compute_logprobs(logits)
+            return gather_logprobs(logprobs, next_tokens, max_logprobs)
+
+        logprobs = jit_compute_and_gather_logprobs(logits, next_tokens,
+                                                   max_logprobs)
+        return jax.tree.map(lambda x: j2t(x.astype(jnp.float32)), logprobs)
 
     def _prepare_dp_input_metadata(self,
                                    scheduler_output: "VllmSchedulerOutput"):