diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py
index 5de06a828b..c4c7d9f225 100644
--- a/src/maxtext/inference/maxengine/maxengine.py
+++ b/src/maxtext/inference/maxengine/maxengine.py
@@ -169,6 +169,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)
@@ -1983,11 +1985,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):
@@ -2011,11 +2017,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