(path, partial_cache, full_cache, annotations)
| 1075 | unboxed_prefix["cache"] = self._maybe_unstack_prefill_result_cache(unboxed_prefix["cache"]) |
| 1076 | |
| 1077 | def copy(path, partial_cache, full_cache, annotations): |
| 1078 | path_key = path[-1].key |
| 1079 | if path_key in [ |
| 1080 | "cache_ar_index", |
| 1081 | "cached_ar_key", |
| 1082 | "cached_ar_value", |
| 1083 | "cached_ar_key_scale", |
| 1084 | "cached_ar_value_scale", |
| 1085 | ]: |
| 1086 | return full_cache # we don't even zero these out because we can mask them out. |
| 1087 | |
| 1088 | batch_idx = -1 |
| 1089 | if "cache_batch" in annotations: |
| 1090 | batch_idx = annotations.index("cache_batch") |
| 1091 | elif "cache_scale_batch" in annotations: |
| 1092 | batch_idx = annotations.index("cache_scale_batch") |
| 1093 | |
| 1094 | if batch_idx < 0: |
| 1095 | raise ValueError(f"Batch index {batch_idx=} shouldn't be less than zero for {path_key}, got {annotations=}") |
| 1096 | |
| 1097 | for slot in slots: |
| 1098 | if path_key == "cache_ar_segment_id": |
| 1099 | ### goal: zero this out in case there is existing data |
| 1100 | s = list(full_cache.shape) |
| 1101 | s[batch_idx] = 1 |
| 1102 | zeros = jnp.zeros(tuple(s), dtype=jnp.int32) |
| 1103 | full_cache = jax.lax.dynamic_update_index_in_dim(full_cache, zeros, slot, batch_idx) |
| 1104 | elif path_key == "cache_prefill_segment_id": |
| 1105 | s = list(full_cache.shape) |
| 1106 | s[batch_idx] = 1 |
| 1107 | zeros = jnp.zeros(tuple(s), dtype=jnp.int32) |
| 1108 | ## zero out in case prefill cache is too small to cover |
| 1109 | full_cache = jax.lax.dynamic_update_index_in_dim(full_cache, zeros, slot, batch_idx) |
| 1110 | ## copy prefill cachce |
| 1111 | full_cache = jax.lax.dynamic_update_index_in_dim(full_cache, partial_cache, slot, batch_idx) |
| 1112 | elif path_key == "cached_ar_lengths": |
| 1113 | full_cache = full_cache.at[slot].set(0) |
| 1114 | elif path_key in [ |
| 1115 | "cached_prefill_key", |
| 1116 | "cached_prefill_value", |
| 1117 | "cached_prefill_key_scale", |
| 1118 | "cached_prefill_value_scale", |
| 1119 | ]: |
| 1120 | full_cache = jax.lax.dynamic_update_index_in_dim(full_cache, partial_cache, slot, batch_idx) |
| 1121 | else: |
| 1122 | raise ValueError(f"We don't have a strategy for inserting {path_key}") |
| 1123 | |
| 1124 | return full_cache |
| 1125 | |
| 1126 | inserted_cache = jax.tree_util.tree_map_with_path( |
| 1127 | copy, |
no outgoing calls