Non-JIT wrapper for inserting prefill cache.
(
self,
prefix: Prefix,
decode_state: DecodeState,
slot: int,
request_id: uuid.UUID | None = None,
)
| 1303 | } |
| 1304 | |
| 1305 | def insert( |
| 1306 | self, |
| 1307 | prefix: Prefix, |
| 1308 | decode_state: DecodeState, |
| 1309 | slot: int, |
| 1310 | request_id: uuid.UUID | None = None, |
| 1311 | ) -> DecodeState: |
| 1312 | """Non-JIT wrapper for inserting prefill cache.""" |
| 1313 | |
| 1314 | current_page_state = None |
| 1315 | if self.config.attention == "paged" and self.page_manager is not None: |
| 1316 | if self.page_state is None: |
| 1317 | self.page_state = self.page_manager.get_initial_page_state() |
| 1318 | current_page_state = self.page_state |
| 1319 | |
| 1320 | updated_decode_state = self._insert_jit( |
| 1321 | prefix=prefix, |
| 1322 | decode_state=decode_state, |
| 1323 | slot=slot, |
| 1324 | page_state_in=current_page_state, |
| 1325 | ) |
| 1326 | |
| 1327 | # Update the PageState after the JIT call |
| 1328 | if self.config.attention == "paged" and self.page_manager is not None and self.page_state is not None: |
| 1329 | new_has_active_page = self.page_state.has_active_page.at[slot].set(True) |
| 1330 | self.page_state = self.page_state.replace(has_active_page=new_has_active_page) |
| 1331 | return updated_decode_state |
| 1332 | |
| 1333 | @functools.partial( |
| 1334 | jax.jit, |