Cache compiled JAX init HLOs in MaxEngine.

src/maxtext/inference/maxengine/maxengine.py

@@ -164,6 +164,8 @@ def __init__(self, config: Any, devices: Any | None = None):
     self.decode_state_layouts = None
     self.param_layouts = None
     self.rng = None
+    self._compiled_initialize_fn = None
+    self._compiled_init_cache_fn = None
 
   def print_stats(self, label: str):
     max_utils.print_mem_stats(label)
@@ -1974,11 +1976,15 @@ def init(abstract_params):
         mesh_annotations,
     )
 
-    @functools.partial(jax.jit, out_shardings=shardings)
-    def initialize():
-      return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs)
+    if self._compiled_initialize_fn is None:
 
-    init_state = initialize()
+      @functools.partial(jax.jit, out_shardings=shardings)
+      def initialize():
+        return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs)
+
+      self._compiled_initialize_fn = initialize
+
+    init_state = self._compiled_initialize_fn()
     cache = init_state["cache"]
 
     def is_lp(k):
@@ -2002,11 +2008,15 @@ def _init_decode_state_nnx(self) -> DecodeState:
     # AR-mode cache so the batch dim matches generate's input shape.
     cache_dict_abs = self._nnx_init_cache_dict(mode=MODEL_MODE_AUTOREGRESSIVE)
 
-    @functools.partial(jax.jit, out_shardings=(self.kv_cache_shardings,))
-    def _init_cache():
-      return (jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict_abs),)
+    if self._compiled_init_cache_fn is None:
+
+      @functools.partial(jax.jit, out_shardings=(self.kv_cache_shardings,))
+      def _init_cache():
+        return (jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict_abs),)
+
+      self._compiled_init_cache_fn = _init_cache
 
-    (cache,) = _init_cache()
+    (cache,) = self._compiled_init_cache_fn()
 
     # Per-leaf logical axes for bulk_insert's "cache_batch" lookup. Use model_ar
     # so segment_id leaves carry CACHE_BATCH (under PREFILL they'd carry