Put data transfer outside of jitted functions

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

Author: Lumosis

tpu_inference/runner/compilation_manager.py

@@ -483,11 +483,8 @@ def _precompile_eagle3_helpers(self) -> None:
         target_hidden_state_loop = self._create_dummy_tensor(
             (self.runner.max_num_reqs, hidden_size), dtype,
             NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
-        new_seq_lens_cpu = np.ones((self.runner.max_num_reqs, ), jnp.int32)
         next_token_ids = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
-        new_query_start_loc_cpu = np.ones((self.runner.max_num_reqs + 1, ),
-                                          jnp.int32)
         last_token_indices = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
         for num_tokens in self.runner.num_tokens_paddings:
@@ -535,6 +532,7 @@ def prepare_inputs_in_jit_wrapper(
                 return target_hidden_states, input_ids, last_token_indices
 
             token_indices_cpu = np.ones((num_tokens, ), dtype=np.int32)
+            token_indices = jnp.asarray(token_indices_cpu, dtype=jnp.int32)
             input_ids = self._create_dummy_tensor(
                 (num_tokens, ), jnp.int32,
                 NamedSharding(self.runner.mesh, PartitionSpec()))
@@ -555,9 +553,9 @@ def prepare_inputs_in_jit_wrapper(
             self._run_compilation(
                 "eagle3_prepare_inputs_in_jit",
                 prepare_inputs_in_jit_wrapper,
-                token_indices_cpu,
-                new_query_start_loc_cpu,
-                new_seq_lens_cpu,
+                token_indices,
+                query_start_loc,
+                seq_lens,
                 input_ids,
                 aux_hidden_states,
                 attention_metadata,
@@ -630,13 +628,6 @@ def draft_model_fn_wrapper(
             hidden_states = self._create_dummy_tensor(
                 (num_tokens, hidden_size), jnp.bfloat16,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
-            self._run_compilation(
-                "eagle3_select_draft_token_ids_in_jit",
-                self.runner.drafter._select_draft_token_ids_in_jit,
-                hidden_states,
-                last_token_indices,
-                num_tokens=num_tokens,
-            )
 
             self._run_compilation(
                 "eagle3_select_and_stack_draft_token_ids_in_jit",
@@ -648,9 +639,10 @@ def draft_model_fn_wrapper(
 
             self._run_compilation(
                 "eagle3_select_positions_and_hidden_states_in_jit",
-                self.runner.drafter._select_positions_and_hidden_states_in_jit,
+                self.runner.drafter._select_inputs_for_loop_in_jit,
                 positions,
                 hidden_states,
+                hidden_states,
                 last_token_indices,
                 num_tokens=num_tokens,
             )

tpu_inference/spec_decode/jax/eagle3.py

@@ -217,30 +217,30 @@ def prepare_inputs(
         # Update seq_lens for active requests only: new_seq_lens = s - n.
         new_seq_lens_cpu = seq_lens_cpu - nrt_cpu
 
-        return self._prepare_inputs_in_jit(token_indices_cpu,
-                                           new_query_start_loc_cpu,
-                                           new_seq_lens_cpu, input_ids,
+        token_indices = jnp.asarray(token_indices_cpu, dtype=jnp.int32)
+        query_start_loc, seq_lens = device_array(self.mesh, (
+            new_query_start_loc_cpu,
+            new_seq_lens_cpu,
+        ))
+
+        return self._prepare_inputs_in_jit(token_indices, query_start_loc,
+                                           seq_lens, input_ids,
                                            aux_hidden_states, attn_metadata,
                                            next_token_ids, block_tables)
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _prepare_inputs_in_jit(
         self,
-        token_indices_cpu: np.ndarray,
-        new_query_start_loc_cpu: np.ndarray,
-        new_seq_lens_cpu: np.ndarray,
+        token_indices: jax.Array,
+        query_start_loc: jax.Array,
+        seq_lens: jax.Array,
         input_ids: jax.Array,
         aux_hidden_states: tuple[jax.Array, ...],
         attn_metadata: AttentionMetadata,
         next_token_ids: jax.Array,
         block_tables: jax.Array,
     ) -> tuple[jax.Array, jax.Array, jax.Array, AttentionMetadata]:
 
-        token_indices = jnp.asarray(token_indices_cpu, dtype=jnp.int32)
-        query_start_loc, seq_lens = device_array(self.mesh, (
-            new_query_start_loc_cpu,
-            new_seq_lens_cpu,
-        ))
         # Select tokens and hidden states.
         target_token_ids = input_ids[token_indices]
         target_hidden_states = jnp.concatenate(
@@ -259,7 +259,7 @@ def _prepare_inputs_in_jit(
             request_distribution=attn_metadata.request_distribution,
         )
 
-        target_hidden_states, input_ids, last_token_indices, block_tables = self._prepare_draft_inputs(
+        target_hidden_states, input_ids, last_token_indices, block_tables = self._prepare_draft_inputs_in_jit(
             target_hidden_states, query_start_loc, target_token_ids,
             next_token_ids, block_tables)
 
@@ -272,16 +272,6 @@ def _prepare_draft_inputs_in_jit(
         target_token_ids: jax.Array, next_token_ids: jax.Array,
         block_tables: jax.Array
     ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
-        return self._prepare_draft_inputs(target_hidden_states,
-                                          query_start_loc, target_token_ids,
-                                          next_token_ids, block_tables)
-
-    def _prepare_draft_inputs(
-        self, target_hidden_states: jax.Array, query_start_loc: jax.Array,
-        target_token_ids: jax.Array, next_token_ids: jax.Array,
-        block_tables: jax.Array
-    ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
-
         target_hidden_states = self.combine_hidden_states_fn(
             self.state, target_hidden_states)
 
@@ -294,13 +284,13 @@ def _prepare_draft_inputs(
 
         return target_hidden_states, input_ids, last_token_indices, block_tables
 
-    @functools.partial(jax.jit, static_argnums=(0, ))
-    def _select_draft_token_ids_in_jit(
+    def _select_draft_token_ids(
         self,
         hidden_states: jax.Array,
         last_token_indices: jax.Array,
     ) -> jax.Array:
-        return self._select_draft_token_ids(hidden_states, last_token_indices)
+        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(
@@ -312,30 +302,22 @@ def _select_and_stack_draft_token_ids_in_jit(
                                                        last_token_indices)
         return jnp.stack([draft_token_ids], axis=1)
 
-    def _select_draft_token_ids(
-        self,
-        hidden_states: jax.Array,
-        last_token_indices: jax.Array,
-    ) -> jax.Array:
-        sample_hidden_states = hidden_states[last_token_indices]
-        return self._get_draft_token_ids(sample_hidden_states)
-
-    def _get_draft_token_ids(self, hidden_states: jax.Array) -> jax.Array:
+    @functools.partial(jax.jit, static_argnums=(0, ))
+    def _get_draft_token_ids_in_jit(self,
+                                    hidden_states: jax.Array) -> jax.Array:
         lora_metadata = None
         logits = self.compute_logits_fn(self.state, hidden_states,
                                         lora_metadata)
         return jnp.argmax(logits, axis=-1)
 
     @functools.partial(jax.jit, static_argnums=(0, ))
-    def _get_draft_token_ids_in_jit(self,
-                                    hidden_states: jax.Array) -> jax.Array:
-        return self._get_draft_token_ids(hidden_states)
-
-    @functools.partial(jax.jit, static_argnums=(0, ))
-    def _select_positions_and_hidden_states_in_jit(
-            self, positions: jax.Array, hidden_states: jax.Array,
+    def _select_inputs_for_loop_in_jit(
+            self, positions: jax.Array, residual: jax.Array,
+            hidden_states: jax.Array,
             last_token_indices: jax.Array) -> tuple[jax.Array, jax.Array]:
-        return positions[last_token_indices], hidden_states[last_token_indices]
+        return positions[last_token_indices], residual[
+            last_token_indices], self._select_draft_token_ids(
+                hidden_states, last_token_indices)
 
     def propose(
         self,
@@ -364,14 +346,12 @@ def propose(
             return kv_caches, self._select_and_stack_draft_token_ids_in_jit(
                 hidden_states, last_token_indices)
 
-        draft_token_ids = self._select_draft_token_ids_in_jit(
-            hidden_states, last_token_indices)
+        positions, hidden_states, draft_token_ids = self._select_inputs_for_loop_in_jit(
+            attn_metadata.input_positions, residual[0], hidden_states,
+            last_token_indices)
 
         draft_token_ids_list = [draft_token_ids]
 
-        positions, hidden_states = self._select_positions_and_hidden_states_in_jit(
-            attn_metadata.input_positions, residual[0], last_token_indices)
-
         for _ in range(self.num_speculative_tokens - 1):
             input_ids_loop = draft_token_ids_list[-1]