Revert previous API changes due to upstream change. (#1155)

Signed-off-by: Lihao Ran <imlihao.ran@gmail.com>

Author: Lumosis

tests/layers/jax/sample/test_rejection_sampler.py

@@ -436,9 +436,6 @@ def run_rejection_sampler_test(
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         assert parsed_output == test_case.expected, \
             f"Test '{test_case.name}': Expected {test_case.expected}, got {parsed_output}"
 
@@ -515,9 +512,6 @@ def test_parse_output_basic(self, rejection_sampler):
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         expected = [[10, 20, 30, 40], [50, 60, 70]]
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
 
@@ -541,9 +535,6 @@ def test_parse_output_with_placeholders(self, rejection_sampler):
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         expected = [[10], [20, 30, 40]]
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
 
@@ -565,9 +556,6 @@ def test_parse_output_invalid_tokens(self, rejection_sampler):
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         expected = [[10, 20]]  # Invalid tokens filtered out
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
 
@@ -589,9 +577,6 @@ def test_parse_output_empty_sequences(self, rejection_sampler):
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         expected = [[50], [60]]  # Only bonus tokens
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
 
@@ -647,9 +632,6 @@ def test_extreme_padding(self, rejection_sampler, test_helper):
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         expected = [[1, 5]]  # Should ignore all padding
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
 
@@ -795,9 +777,6 @@ def test_single_long_sequence(self, rejection_sampler, test_helper):
             batch_size=len(num_draft_tokens),
             padded_tokens_length=int(sum(num_draft_tokens)))
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         expected = [list(range(1, 28)) + [99]]  # Tokens up to mismatch point
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
 
@@ -905,9 +884,6 @@ def test_non_greedy_deterministic_with_seed(self, rejection_sampler,
                 num_draft_tokens_cpu=np.asarray(num_draft_tokens),
                 batch_size=1,
                 padded_tokens_length=4)
-
-            # Convert numpy arrays to lists for comparison
-            parsed_output = [x.tolist() for x in parsed_output]
             outputs.append(parsed_output)
 
         # All outputs should be identical with same seed
@@ -1088,9 +1064,6 @@ def test_non_greedy_empty_sequence(self, rejection_sampler, test_helper):
             batch_size=2,
             padded_tokens_length=0)
 
-        # Convert numpy arrays to lists for comparison
-        parsed_output = [x.tolist() for x in parsed_output]
-
         # Should get bonus tokens for empty sequences
         expected = [[77], [88]]
         assert parsed_output == expected, f"Expected {expected}, got {parsed_output}"
@@ -1179,10 +1152,6 @@ def test_non_greedy_vs_greedy_same_perfect_case(self, rejection_sampler,
         non_greedy_parsed = rejection_sampler.parse_output(
             non_greedy_output, VOCAB_SIZE, np.asarray(num_draft_tokens), 1, 3)
 
-        # Convert numpy arrays to lists for comparison
-        greedy_parsed = [x.tolist() for x in greedy_parsed]
-        non_greedy_parsed = [x.tolist() for x in non_greedy_parsed]
-
         # For perfect match, greedy should have all tokens + bonus
         assert greedy_parsed == [[5, 15, 25, 99]]
 

tests/runner/test_speculative_decoding_manager.py

@@ -321,7 +321,7 @@ def test_propose_eagle3_draft_token_ids(self,
         )
 
         # Inputs
-        sampled_token_ids = [np.array([1]), np.array([2])]
+        sampled_token_ids = [[1], [2]]
         aux_hidden_states = MagicMock()
         attn_metadata = MagicMock()
         attn_metadata.seq_lens.shape = [2]

tpu_inference/layers/jax/sample/rejection_sampler.py

@@ -128,7 +128,7 @@ def parse_output(
         num_draft_tokens_cpu: np.ndarray,
         batch_size: int,
         padded_tokens_length: int,
-    ) -> list[np.ndarray]:
+    ) -> list[list[int]]:
         """Parse the output of the rejection sampler.
 
         Args:
@@ -177,7 +177,7 @@ def parse_output(
             else:
                 seq_tokens = valid_main_tokens
 
-            outputs.append(seq_tokens)
+            outputs.append(seq_tokens.tolist())
             start_idx = end_idx
 
         return outputs

tpu_inference/runner/speculative_decoding_manager.py

@@ -78,7 +78,7 @@ def propose_draft_token_ids(
 
     def propose_eagle3_draft_token_ids(
         self,
-        sampled_token_ids: list[np.ndarray],
+        sampled_token_ids: list[list[int]],
         aux_hidden_states: Optional[tuple[jnp.ndarray, ...]],
         attn_metadata: AttentionMetadata,
         spec_decode_metadata: Optional[SpecDecodeMetadata],
@@ -91,7 +91,7 @@ def propose_eagle3_draft_token_ids(
         req_ids = self.runner.input_batch.req_ids
         next_token_ids: list[int] = []
         for i, token_ids in enumerate(sampled_token_ids):
-            if token_ids.size != 0:
+            if token_ids:
                 # Common case.
                 next_token_id = token_ids[-1]
             else:

tpu_inference/runner/tpu_runner.py

@@ -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,
-                             LogprobsTensors, ModelRunnerOutput)
+                             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,10 +122,9 @@ 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 = [token_id for token_id in selected_token_ids]
+        valid_sampled_token_ids = selected_token_ids.tolist()
         for i in self._discard_sampled_tokens_req_indices:
-            valid_sampled_token_ids[i] = np.array([])
+            valid_sampled_token_ids[i].clear()
         self._model_runner_output.sampled_token_ids = valid_sampled_token_ids
         return self._model_runner_output
 
@@ -614,11 +613,11 @@ 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 = [token_id for token_id in selected_token_ids]
+        valid_sampled_token_ids = selected_token_ids.tolist()
 
         # Mask out the sampled tokens that should not be sampled.
         for i in pre_discard_sampled_tokens_req_indices:
-            valid_sampled_token_ids[i] = np.array([])
+            valid_sampled_token_ids[i].clear()
         # Append sampled tokens
         for pre_req_idx, req_state, _ in pre_request_seq_lens:
             sampled_ids = valid_sampled_token_ids[pre_req_idx]
@@ -940,9 +939,7 @@ 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 = [
-                token_id for token_id in selected_token_ids
-            ]
+            valid_sampled_token_ids = selected_token_ids.tolist()
         else:
             valid_sampled_token_ids = self.rejection_sampler.parse_output(
                 next_tokens, self.input_batch.vocab_size,
@@ -951,11 +948,11 @@ def _sample_from_logits(
 
         # Mask out the sampled tokens that should not be sampled.
         for i in discard_sampled_tokens_req_indices:
-            valid_sampled_token_ids[i] = np.array([])
+            valid_sampled_token_ids[i].clear()
         # Append sampled tokens
         for req_idx, req_state, _ in request_seq_lens:
             sampled_ids = valid_sampled_token_ids[req_idx]
-            if sampled_ids.size == 0:
+            if not sampled_ids:
                 continue
 
             start_idx = self.input_batch.num_tokens_no_spec[req_idx]
@@ -1018,8 +1015,7 @@ def select_local_fn(local_array, local_indices):
 
     @staticmethod
     @functools.partial(jax.jit, static_argnames=("max_logprobs", ))
-    def _compute_and_gather_logprobs(logits, next_tokens,
-                                     max_logprobs) -> LogprobsTensors:
+    def _compute_and_gather_logprobs(logits, next_tokens, max_logprobs):
         logprobs = compute_logprobs(logits)
         return gather_logprobs(logprobs, next_tokens, max_logprobs)