Remove cached_ar_key/cached_ar_value in prefill results to save memory.

MaxText/maxengine.py

@@ -84,7 +84,8 @@ def load_params(self, *args, **kwargs) -> Params:
     )
     self.kv_cache_annotations = max_utils.get_kv_cache_annotations(self.model, self.config, self.rng, self._mesh)
     self.kv_cache_shardings = jax.tree_util.tree_map(
-      lambda x: jax.sharding.NamedSharding(self._mesh, x), self.kv_cache_annotations)
+        lambda x: jax.sharding.NamedSharding(self._mesh, x), self.kv_cache_annotations
+    )
 
     if not self.model.quant:
       self.abstract_params = jax.tree_util.tree_map(
@@ -171,6 +172,18 @@ def prefill(
         flat_logits, (0, true_length - 1, 0), (flat_logits.shape[0], 1, flat_logits.shape[2])
     )
     selected_logits = jax.lax.with_sharding_constraint(selected_logits, self.replicated_sharding)
+
+    # delete the "cached_ar_key" and "cached_ar_value" keys would cause error in insert() function
+    # so set cached_ar_key/cached_ar_value to dummy shape (1,) value to save memory
+    for i in range(self.config.num_decoder_layers):
+      if f"layers_{i}" in new_vars["cache"]["decoder"]:
+        new_vars["cache"]["decoder"][f"layers_{i}"]["self_attention"]["AttentionOp_0"]["cached_ar_key"] = jnp.zeros(
+            (1), dtype=jnp.int32
+        )
+        new_vars["cache"]["decoder"][f"layers_{i}"]["self_attention"]["AttentionOp_0"]["cached_ar_value"] = jnp.zeros(
+            (1), dtype=jnp.int32
+        )
+
     return {
         "logits": selected_logits,
         "cache": new_vars["cache"],