[Spec Decoding] Fix precompilation (#960)

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

Author: Lumosis

tests/runner/test_speculative_decoding_manager.py

@@ -302,14 +302,15 @@ def test_propose_eagle3_draft_token_ids(self,
         # Mock drafter methods
         mock_attn_metadata = MagicMock()
         mock_target_token_ids = MagicMock()
+        mock_last_token_indices = MagicMock()
         mock_target_hidden_states = MagicMock()
         self.runner.drafter.prepare_inputs.return_value = (
-            mock_attn_metadata,
-            mock_target_token_ids,
             mock_target_hidden_states,
+            mock_target_token_ids,
+            mock_last_token_indices,
+            mock_attn_metadata,
         )
-        mock_draft_token_ids = MagicMock()
-        mock_draft_token_ids.tolist.return_value = [[10, 11], [20, 21]]
+        mock_draft_token_ids = [[10, 11], [20, 21]]
         self.runner.drafter.propose.return_value = (
             self.runner.kv_caches,
             mock_draft_token_ids,

tests/spec_decode/test_eagle3.py

@@ -86,7 +86,7 @@ def test_prepare_inputs():
     proposer.state = None  # Mock state
     proposer.runner.input_batch.block_table = [mock.MagicMock()]
     # Mock the block table return value (2D array)
-    (proposer.runner.input_batch.block_table[0].get_device_tensor.return_value
+    (proposer.runner.input_batch.block_table[0].get_cpu_tensor.return_value
      ) = jnp.zeros((num_reqs, max_num_blocks_per_req), dtype=jnp.int32)
 
     # --- Setup sequence data ---
@@ -289,6 +289,8 @@ def mock_combine_hidden_states_fn(state, hidden_states):
         target_hidden_states,
     )
 
+    if draft_token_ids.ndim == 1:
+        draft_token_ids = jnp.expand_dims(draft_token_ids, axis=-1)
     # Assertions
     assert draft_token_ids.shape == (batch_size, num_speculative_tokens)
 

tpu_inference/runner/compilation_manager.py

@@ -428,9 +428,7 @@ def _precompile_eagle3_helpers(self) -> None:
         draft_kv_cache_group_id = num_kv_cache_groups - 1
         block_tables = self.runner.input_batch.block_table[
             draft_kv_cache_group_id].get_cpu_tensor().reshape(-1)
-        block_tables_first_spec = jax.device_put(
-            block_tables, NamedSharding(self.runner.mesh, PartitionSpec()))
-        block_tables_loop = jax.device_put(
+        block_tables = jax.device_put(
             block_tables, NamedSharding(self.runner.mesh,
                                         PartitionSpec(None, )))
 
@@ -447,7 +445,7 @@ def _precompile_eagle3_helpers(self) -> None:
         self._run_compilation(
             "_update_inputs_for_loop_speculation for the subsequent loops",
             self.runner.drafter._update_inputs_for_loop_speculation,
-            selected_positions, seq_lens, block_tables_loop)
+            selected_positions, seq_lens, block_tables)
 
         request_distribution = np.array([0, 0, 0], dtype=np.int32)
         request_distribution = device_array(self.runner.mesh,
@@ -498,7 +496,7 @@ def _precompile_eagle3_helpers(self) -> None:
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32)
             attention_metadata = AttentionMetadata(
                 input_positions=positions,
-                block_tables=block_tables_first_spec,
+                block_tables=block_tables,
                 seq_lens=seq_lens,
                 query_start_loc=query_start_loc,
                 request_distribution=request_distribution,
@@ -520,11 +518,7 @@ def filter_token_and_prepare_initial_inputs_wrapper(
                     num_reqs)
                 return target_hidden_states, input_ids, last_token_indices
 
-            token_indices = self._create_dummy_tensor((num_tokens, ),
-                                                      jnp.int32)
-            input_ids = self._create_dummy_tensor(
-                (num_tokens, ), jnp.int32,
-                NamedSharding(self.runner.mesh, PartitionSpec()))
+            input_ids = self._create_dummy_tensor((num_tokens, ), jnp.int32)
             aux_hidden_states = [
                 self._create_dummy_tensor(
                     (num_tokens, hidden_size), jnp.bfloat16,
@@ -539,22 +533,29 @@ def filter_token_and_prepare_initial_inputs_wrapper(
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
             ]
-            self._run_compilation(
-                "eagle3_filter_token_and_prepare_initial_inputs",
-                filter_token_and_prepare_initial_inputs_wrapper,
-                token_indices,
-                query_start_loc,
-                seq_lens,
-                input_ids,
-                aux_hidden_states,
-                attention_metadata,
-                next_token_ids,
-                device_array(
-                    self.runner.mesh,
-                    np.asarray([self.runner.input_batch.num_reqs],
-                               dtype=jnp.int32)),
-                num_tokens=num_tokens,
-            )
+            # TODO(ranlihao): This will increase the precompilation latency. Find proper range for token_indices.
+            for padded_total_num_tokens in [
+                    num_tokens,
+                    min(num_tokens * 2, self.runner.num_tokens_paddings[-1])
+            ]:
+                token_indices = self._create_dummy_tensor(
+                    (padded_total_num_tokens, ), jnp.int32)
+                self._run_compilation(
+                    "eagle3_filter_token_and_prepare_initial_inputs",
+                    filter_token_and_prepare_initial_inputs_wrapper,
+                    token_indices,
+                    query_start_loc,
+                    seq_lens,
+                    input_ids,
+                    aux_hidden_states,
+                    attention_metadata,
+                    next_token_ids,
+                    device_array(
+                        self.runner.mesh,
+                        np.asarray([self.runner.input_batch.num_reqs],
+                                   dtype=jnp.int32)),
+                    num_tokens=num_tokens,
+                )
 
             def draft_model_fn_wrapper(
                 state,
@@ -572,6 +573,9 @@ def draft_model_fn_wrapper(
             target_hidden_states = self._create_dummy_tensor(
                 (num_tokens, hidden_size), dtype,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, "model")))
+            input_ids = self._create_dummy_tensor(
+                (num_tokens, ), jnp.int32,
+                NamedSharding(self.runner.mesh, PartitionSpec()))
             self._run_compilation(
                 "eagle3_draft_model_fn",
                 draft_model_fn_wrapper,
@@ -602,7 +606,6 @@ def draft_model_fn_wrapper(
             attention_metadata.query_start_loc = jax.device_put(
                 attention_metadata.query_start_loc,
                 NamedSharding(self.runner.mesh, PartitionSpec()))
-            attention_metadata.block_tables = block_tables_loop
             attention_metadata.input_positions = self._create_dummy_tensor(
                 (self.runner.max_num_reqs, ), jnp.int32)
             self._run_compilation(