hacky way to test aot in jetstream

Author: jwyang-google

jetstream/core/orchestrator.py

@@ -102,7 +102,11 @@
 from jetstream.core.metrics.prometheus import JetstreamMetricsCollector
 import numpy as np
 
-log_level = os.getenv("LOG_LEVEL", "WARNING").upper()
+from jax.experimental import layout as jax_layout
+DLL = jax_layout.DeviceLocalLayout
+Layout = jax_layout.Layout
+
+log_level = os.getenv("LOG_LEVEL", "DEBUG").upper()
 
 logger = logging.getLogger("JetstreamLogger")
 logger.propagate = False
@@ -405,6 +409,63 @@ def __init__(
 
     self._jax_padding = jax_padding
 
+    ##### hacky code using auto layout for interleaved engine
+    self.engine = self._generate_engines[0]
+    self.params = self._generate_params[0]
+    logger.debug("Compiling generate function")
+    self._generate_executable, self.params, self._decode_state_executable = self.engine.aot_compile(
+        self.params, pass_rng_shape=False
+    )
+    self.decode_state = self._decode_state_executable(None)
+
+    # prefill
+    interesting_buckets = [
+        64,
+        128,
+        256,
+        512,
+        1024,
+    ]
+
+    self._cached_prefill = {}
+    self._cached_insert = {}
+    for length in interesting_buckets:
+      i32_scalar = jax.ShapeDtypeStruct((), int)
+      logger.debug("Compiling prefill: %d", length)
+      input_data = jax.ShapeDtypeStruct((length,), jax.numpy.dtype("int32"))
+
+      self._cached_prefill[length] = (
+          jax.jit(
+            self.engine.prefill_aot,
+            in_shardings=(self.engine.param_layouts, None, None),
+            out_shardings=(Layout(DLL.AUTO), Layout(DLL.AUTO)),
+          ).lower(self.params, input_data, i32_scalar)
+      ).compile(compiler_options=None)
+
+      logger.debug("Generate dummy prefix: %d", length)
+      dummy_tokens = jax.numpy.ones(shape=(length,), dtype=jax.numpy.dtype("int32"))
+      prefix_shapes = jax.eval_shape(self.engine.prefill_aot, self.params, dummy_tokens, 1)
+      
+      logger.debug("Compiling insert: %d", length)
+      prefill_output_layout, _ = self._cached_prefill[length].output_layouts
+      logger.debug("Prefill output layout: {}".format(prefill_output_layout))
+      logger.debug("Prefix shapes: {}".format(prefix_shapes))
+      i32_scalar = jax.ShapeDtypeStruct((), int)
+      self._cached_insert[length] = (
+          jax.jit(
+            self.engine.insert,
+            in_shardings=(prefill_output_layout, self.engine.decode_state_layouts, None),
+            out_shardings=(self.engine.decode_state_layouts),
+            donate_argnames=("decode_state"),
+          ).lower(prefix_shapes[0], self.engine.decode_state_shapes, i32_scalar)
+      ).compile(compiler_options=None)
+
+    self._prefill_engines[0] = self.engine
+    self._generate_engines[0] = self.engine
+    self._prefill_params[0] = self.params
+    self._generate_params[0] = self.params
+
+
     # Create all threads
     self._prefill_threads = [
         JetThread(
@@ -759,10 +820,11 @@ def _prefill_thread(self, idx: int):
           )
         else:
           # Compute new kv cache for the prefill_content.
-          prefill_result, first_token = prefill_engine.prefill(
-              params=final_prefill_params,
-              padded_tokens=padded_tokens,
-              true_length=true_length,
+          assert padded_tokens.shape[0] in self._cached_prefill
+          prefill_result, first_token = self._cached_prefill[padded_tokens.shape[0]](
+              final_prefill_params,
+              padded_tokens,
+              true_length,
           )
 
         request.complete = np.zeros(
@@ -967,10 +1029,11 @@ def _insert_if_possible(
         else:
           break
 
-      decode_state = generate_engine.insert(
+      length = new_request.prefill_result['cache']['decoder']['layers_0']['self_attention']['KVCache_0']['cache_prefill_segment_id'].value.shape[1]
+      decode_state = self._cached_insert[length](
           new_request.prefill_result,
           decode_state,
-          slot=slot,
+          slot,
           # request_id=new_request.request_id,
       )
       ThreadDebugLog(
@@ -1115,9 +1178,9 @@ def _generate_thread(self, idx: int):
     # Keep track of what step tokens were generated at.
     generate_timestep = 0
     # State to store things like running kv cache in.
-    decode_state = generate_engine.init_decode_state()
-
+    decode_state = self.decode_state
     generate_params = self._generate_params[idx]
+
     thread_name = f"Generate thread {idx}"
     ThreadDebugLog(thread_name, f"Generate params {idx} loaded.")
     time_of_last_generate = time.time()
@@ -1178,8 +1241,8 @@ def _generate_thread(self, idx: int):
       ), "At this point we must have some requests inserted into the slots."
 
       # Now we actually take a generate step on requests in the slots.
-      decode_state, sampled_tokens = generate_engine.generate(
-          generate_params, decode_state
+      decode_state, sampled_tokens = self._generate_executable(
+          generate_params, decode_state, None
       )
       sampled_tokens.copy_to_host_async()
       # Respond to detokenization backpressure.