Reuse KV cache of prefixes #5572
Conversation
| for i in range(num_already_cached_blocks, n_blocks): | ||
| chunk = tokens[:(i + 1) * self.block_size] | ||
| hash = token_ids_to_hash(chunk) | ||
| if hash not in self.tokens_to_blocks: |
There was a problem hiding this comment.
Do we still need this if statement if in the for loop you are already starting from num_already_cached_blocks.
There was a problem hiding this comment.
I'm not sure if it is necessary but there might be a complicated case. Assume we are running two requests.
Request 1 is using cached blocks 0, 1, 2, and just generated the last token of the current block. Then it saves the generated sequence and a hash.
Request 2 is also using the cached block 0, 1, 2, but the generation is a few steps later than Request 1. They are not sharing the last block. But the request may generate exact same tokens for the last block. So, it will try to update the cache with the same hash.
In this case, I didn't want to overwrite the cache.
There was a problem hiding this comment.
sounds good, so this is to ensure we are sharing the recently generated blocks.
This PR implements reusing KV cache prefixes.
See microsoft/DeepSpeed-MII#484 for details.