Add warmup + formatting

vllm/worker/tpu_model_runner.py

@@ -14,6 +14,7 @@ from vllm.utils import pad_to_max_length
 
 # DELETE
 from jax_smi import initialise_tracking
+
 initialise_tracking()
 
 logger = init_logger(__name__)
@@ -60,6 +61,51 @@ class TPUModelRunner:
         params = load_and_format_params(model_dir + "/7b/")["transformer"]
         self.params = {"params": params}
 
+    def warmup_model(self, tpu_caches: List[Tuple[jax.Array,
+                                                  jax.Array]]) -> None:
+        # Prefill
+        logger.info("Warming up the model...")
+        start = time.time()
+        for batch_size in [1]:
+            for seq_len in [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]:
+                if batch_size * seq_len > 8192:
+                    continue
+                token_ids = jnp.zeros((batch_size, seq_len), dtype=jnp.int32)
+                position_ids = jnp.zeros((batch_size, seq_len),
+                                         dtype=jnp.int32)
+                slot_mapping = jnp.zeros((batch_size, seq_len),
+                                         dtype=jnp.int32)
+                block_tables = None
+                context_lens = None
+                prompt_lens = jnp.ones((batch_size, ), dtype=jnp.int32)
+
+                # Dummy run.
+                _, tpu_caches = self.compiled_fn(self.params, token_ids,
+                                                 position_ids, slot_mapping,
+                                                 block_tables, context_lens,
+                                                 prompt_lens, tpu_caches)
+        end = time.time()
+        logger.info(f"Prefill warmup done in {(end - start):.2f} seconds.")
+
+        # Decode
+        start = time.time()
+        for batch_size in [1, 2, 4] + [8 * i for i in range(1, 17)]:
+            seq_len = 1
+            token_ids = jnp.zeros((batch_size, seq_len), dtype=jnp.int32)
+            position_ids = jnp.zeros((batch_size, seq_len), dtype=jnp.int32)
+            slot_mapping = jnp.zeros((batch_size, seq_len), dtype=jnp.int32)
+            block_tables = jnp.asarray(self.block_tables[:batch_size],
+                                       dtype=jnp.int32)
+            context_lens = jnp.ones((batch_size, ), dtype=jnp.int32)
+            prompt_lens = jnp.ones((batch_size, ), dtype=jnp.int32)
+
+            _, tpu_caches = self.compiled_fn(self.params, token_ids,
+                                             position_ids, slot_mapping,
+                                             block_tables, context_lens,
+                                             prompt_lens, tpu_caches)
+        end = time.time()
+        logger.info(f"Decode warmup done in {(end - start):.2f} seconds.")
+
     def _prepare_prompt(
         self,
         seq_group_metadata_list: List[SequenceGroupMetadata],
@@ -160,7 +206,8 @@ class TPUModelRunner:
         slot_mapping = jnp.asarray(slot_mapping, dtype=jnp.int32)
         context_lens = jnp.asarray(context_lens, dtype=jnp.int32)
 
-        block_tables = jnp.asarray(self.block_tables[:batch_size], dtype=jnp.int32)
+        block_tables = jnp.asarray(self.block_tables[:batch_size],
+                                   dtype=jnp.int32)
         input_lens = jnp.asarray([1] * batch_size, dtype=jnp.int32)
         return (input_tokens, input_positions, slot_mapping, block_tables,
                 context_lens, input_lens)
@@ -210,25 +257,26 @@ class TPUModelRunner:
     def execute_model(
         self,
         seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
-        kv_caches: List[jax.Array],
+        kv_caches: List[Tuple[jax.Array, jax.Array]],
-    ) -> Tuple[Optional[SamplerOutput], List[jax.Array]]:
+    ) -> Tuple[Optional[SamplerOutput], List[Tuple[jax.Array, jax.Array]]]:
         from vllm.sequence import SequenceOutput, SequenceGroupOutput, Logprob
 
         start = time.time()
         inputs = self.prepare_input_arrays(seq_group_metadata_list)
         end = time.time()
-        print(f"prepare_input_arrays: {(end - start) * 1000:.2f} ms")
+        # print(f"prepare_input_arrays: {(end - start) * 1000:.2f} ms")
 
         start = time.time()
-        next_token_ids, new_kv_caches = self.compiled_fn(self.params, *inputs, kv_caches)
+        next_token_ids, new_kv_caches = self.compiled_fn(
+            self.params, *inputs, kv_caches)
         next_token_ids.block_until_ready()
         end = time.time()
-        print(f"compiled_fn: {(end - start) * 1000:.2f} ms")
+        # print(f"compiled_fn: {(end - start) * 1000:.2f} ms")
 
         start = time.time()
         next_token_ids = jax.device_put(next_token_ids, jax.devices("cpu")[0])
         end = time.time()
-        print(f"jax.device_put: {(end - start) * 1000:.2f} ms")
+        # print(f"jax.device_put: {(end - start) * 1000:.2f} ms")
 
         next_token_ids = next_token_ids.tolist()
         i = 0
@@ -240,7 +288,9 @@ class TPUModelRunner:
             seq_ids = list(seq_group_metadata.seq_data.keys())
             for seq_id in seq_ids:
                 next_token_id = next_token_ids[i]
-                seq_outputs.append(SequenceOutput(seq_id, next_token_id, {next_token_id: Logprob(0.0)}))
+                seq_outputs.append(
+                    SequenceOutput(seq_id, next_token_id,
+                                   {next_token_id: Logprob(0.0)}))
                 i += 1
 
             sampler_outputs.append(SequenceGroupOutput(seq_outputs, None))

vllm/worker/tpu_worker.py

@@ -92,8 +92,7 @@ class TPUWorker(LoraNotSupportedWorkerBase):
         self._warmup_model()
 
     def _warmup_model(self) -> None:
-        # self.model_runner.warmup_model(self.tpu_cache)
+        self.model_runner.warmup_model(self.tpu_cache)
-        pass
 
     def get_cache_block_size_bytes(self) -> int:
         head_size = self.model_config.get_head_size()
@@ -129,8 +128,8 @@ class TPUWorker(LoraNotSupportedWorkerBase):
         if num_seq_groups == 0:
             return {}
 
-        output, kv_caches = self.model_runner.execute_model(seq_group_metadata_list,
+        output, kv_caches = self.model_runner.execute_model(
-                                                 self.tpu_cache)
+            seq_group_metadata_list, self.tpu_cache)
         self.tpu_cache = kv_caches
         return output