Resolve comments

Signed-off-by: Lihao Ran <[email protected]>

Author: Lumosis

tpu_inference/runner/compilation_manager.py

@@ -630,15 +630,7 @@ def draft_model_fn_wrapper(
                 NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
 
             self._run_compilation(
-                "eagle3_select_and_stack_draft_token_ids_in_jit",
-                self.runner.drafter._select_and_stack_draft_token_ids_in_jit,
-                hidden_states,
-                last_token_indices,
-                num_tokens=num_tokens,
-            )
-
-            self._run_compilation(
-                "eagle3_select_positions_and_hidden_states_in_jit",
+                "eagle3_select_inputs_for_loop_in_jit",
                 self.runner.drafter._select_inputs_for_loop_in_jit,
                 positions,
                 hidden_states,

tpu_inference/runner/speculative_decoding_manager.py

@@ -142,8 +142,10 @@ def propose_eagle3_draft_token_ids(
             last_token_indices=last_token_indices,
             target_hidden_states=target_hidden_states,
         )
-        result = draft_token_ids.tolist()
-        return result
+        draft_token_ids = np.array(draft_token_ids)
+        if draft_token_ids.ndim == 1:
+            draft_token_ids = np.expand_dims(draft_token_ids, axis=-1)
+        return draft_token_ids.tolist()
 
     def get_spec_decode_metadata(
         self,

tpu_inference/spec_decode/jax/eagle3.py

@@ -284,6 +284,7 @@ def _prepare_draft_inputs_in_jit(
 
         return target_hidden_states, input_ids, last_token_indices, block_tables
 
+    @functools.partial(jax.jit, static_argnums=(0, ))
     def _select_draft_token_ids(
         self,
         hidden_states: jax.Array,
@@ -292,16 +293,6 @@ def _select_draft_token_ids(
         sample_hidden_states = hidden_states[last_token_indices]
         return self._get_draft_token_ids_in_jit(sample_hidden_states)
 
-    @functools.partial(jax.jit, static_argnums=(0, ))
-    def _select_and_stack_draft_token_ids_in_jit(
-        self,
-        hidden_states: jax.Array,
-        last_token_indices: jax.Array,
-    ) -> jax.Array:
-        draft_token_ids = self._select_draft_token_ids(hidden_states,
-                                                       last_token_indices)
-        return jnp.stack([draft_token_ids], axis=1)
-
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _get_draft_token_ids_in_jit(self,
                                     hidden_states: jax.Array) -> jax.Array:
@@ -343,7 +334,7 @@ def propose(
         )
 
         if self.num_speculative_tokens == 1:
-            return kv_caches, self._select_and_stack_draft_token_ids_in_jit(
+            return kv_caches, self._select_draft_token_ids(
                 hidden_states, last_token_indices)
 
         positions, hidden_states, draft_token_ids = self._select_inputs_for_loop_in_jit(