Retrieve shared KV cache and ids for a placeholder given message context.
(
self,
ph_id: str,
message: str,
)
| 1213 | return new_ph_cache |
| 1214 | |
| 1215 | def fetch_shared_cache( |
| 1216 | self, |
| 1217 | ph_id: str, |
| 1218 | message: str, |
| 1219 | ) -> Tuple[DynamicCache, Dict[str, torch.Tensor], int]: |
| 1220 | """Retrieve shared KV cache and ids for a placeholder given message context.""" |
| 1221 | shared_memory = self.llm._shared_kv_cache_memory |
| 1222 | |
| 1223 | if "user_question" in ph_id: |
| 1224 | return ( |
| 1225 | shared_memory["input"][message][-1], |
| 1226 | shared_memory["input_ids"][message][-1], |
| 1227 | shared_memory["input_drop_num"][message][-1], |
| 1228 | ) |
| 1229 | |
| 1230 | type_str, node_id, *rest = ph_id.split("_") |
| 1231 | is_current = (rest and rest[0] == "current") |
| 1232 | |
| 1233 | key_prefix = "condition" if type_str == "condition" else "response" |
| 1234 | slot_idx = -1 if is_current else -2 |
| 1235 | |
| 1236 | node_memory = shared_memory[node_id] |
| 1237 | |
| 1238 | def _get_slot(bucket_key: str): |
| 1239 | bucket = node_memory.get(bucket_key, {}) |
| 1240 | values = bucket.get(message) |
| 1241 | if not values: |
| 1242 | return None |
| 1243 | try: |
| 1244 | return values[slot_idx] |
| 1245 | except IndexError: |
| 1246 | return None |
| 1247 | |
| 1248 | ph_cache = _get_slot(key_prefix) |
| 1249 | ph_cache_ids = _get_slot(f"{key_prefix}_ids") |
| 1250 | drop_num = _get_slot(f"{key_prefix}_drop_num") |
| 1251 | |
| 1252 | if ph_cache is None: |
| 1253 | raise RuntimeError( |
| 1254 | f"fetch_shared_cache: placeholder {ph_id} for message='{message}' not found." |
| 1255 | ) |
| 1256 | |
| 1257 | return ph_cache, ph_cache_ids, drop_num |
| 1258 | |
| 1259 | @staticmethod |
| 1260 | def trim_token_ids(ids_dict: Dict[str, torch.Tensor], drop_num: int) -> Dict[str, torch.Tensor]: |
no test coverage detected