Insert a single computed prefill cache into KV cache.
(
self,
prefix: Prefix,
decode_state: DecodeState,
slot: int,
request_id: uuid.UUID | None = None, # pylint: disable=unused-argument
page_state_in: PageState | None = None,
)
| 1176 | |
| 1177 | @functools.partial(jax.jit, static_argnums=(0,), donate_argnames=("prefix", "decode_state")) |
| 1178 | def _insert_jit( |
| 1179 | self, |
| 1180 | prefix: Prefix, |
| 1181 | decode_state: DecodeState, |
| 1182 | slot: int, |
| 1183 | request_id: uuid.UUID | None = None, # pylint: disable=unused-argument |
| 1184 | page_state_in: PageState | None = None, |
| 1185 | ) -> DecodeState: |
| 1186 | """Insert a single computed prefill cache into KV cache.""" |
| 1187 | unboxed_prefix = max_utils.unbox_logicallypartioned(prefix) |
| 1188 | unboxed_prefix["cache"] = self._maybe_unstack_prefill_result_cache(unboxed_prefix["cache"]) |
| 1189 | |
| 1190 | def copy(path, partial_cache, full_cache, annotations): |
| 1191 | path_key = path[-1].key |
| 1192 | if path_key in [ |
| 1193 | "cache_ar_index", |
| 1194 | "cached_ar_key", |
| 1195 | "cached_ar_value", |
| 1196 | "cached_ar_key_scale", |
| 1197 | "cached_ar_value_scale", |
| 1198 | ]: |
| 1199 | return full_cache |
| 1200 | |
| 1201 | batch_idx = -1 |
| 1202 | if "cache_batch" in annotations: |
| 1203 | batch_idx = annotations.index("cache_batch") |
| 1204 | elif "cache_scale_batch" in annotations: |
| 1205 | batch_idx = annotations.index("cache_scale_batch") |
| 1206 | |
| 1207 | if batch_idx < 0: |
| 1208 | raise ValueError(f"Batch index {batch_idx=} shouldn't be less than zero for {path_key}, got {annotations=}") |
| 1209 | |
| 1210 | if path_key == "cache_ar_segment_id": |
| 1211 | s = list(full_cache.shape) |
| 1212 | s[batch_idx] = 1 |
| 1213 | zeros = jnp.zeros(tuple(s), dtype=jnp.int32) |
| 1214 | return jax.lax.dynamic_update_index_in_dim(full_cache, zeros, slot, batch_idx) |
| 1215 | elif path_key == "cache_prefill_segment_id": |
| 1216 | s = list(full_cache.shape) |
| 1217 | s[batch_idx] = 1 |
| 1218 | zeros = jnp.zeros(tuple(s), dtype=jnp.int32) |
| 1219 | # zero out in case prefill cache is too small to cover |
| 1220 | full_cache = jax.lax.dynamic_update_index_in_dim(full_cache, zeros, slot, batch_idx) |
| 1221 | # copy prefill cache |
| 1222 | full_cache = jax.lax.dynamic_update_index_in_dim(full_cache, partial_cache, slot, batch_idx) |
| 1223 | return full_cache |
| 1224 | elif path_key == "cached_ar_lengths": |
| 1225 | return full_cache.at[slot].set(0) |
| 1226 | elif path_key in [ |
| 1227 | "cached_prefill_key", |
| 1228 | "cached_prefill_value", |
| 1229 | "cached_prefill_key_scale", |
| 1230 | "cached_prefill_value_scale", |
| 1231 | ]: |
| 1232 | return jax.lax.dynamic_update_index_in_dim(full_cache, partial_cache, slot, batch_idx) |
| 1233 | else: |
| 1234 | raise ValueError(f"We don't have a strategy for inserting {path_key}") |
| 1235 |
no test coverage detected