enable internal kv bucket in llama #720
Conversation
xt574chen
requested review from
ssarkar2,
bhargaveede,
vivekgoe,
mandy-li and
libinta
as code owners
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| if not generation_config.bucket_internal: | ||
| assert generation_config.bucket_size <= 0, "reuse_cache and bucketing flags set together" | ||
| else: | ||
| assert generation_config.bucket_size >= 0, "bucket_internal and bucket_size flags set together" |
There was a problem hiding this comment.
Here, we are in the case where generation_config.bucket_internal is True, so if this assert fails (i.e. generation_config.bucket_size < 0), it means that bucket_size is not set right? But the error message says otherwise
There was a problem hiding this comment.
@puneeshkhanna I see you've corrected some error messages. I hope this update won't cause conflict.
There was a problem hiding this comment.
@xt574chen - I will update my PR once this gets merged first
|
@xt574chen Can you run the following from the root of the repo to make the code style check pass please? |
What does this PR do?
To enhance throughput in scenarios with long new tokens, break down the KV cache into multiples of the bucket width. Use this to compute attention rather than using the entire KV cache.
LLaMA v2 70B (8x, max_input_tokens 128, max_new_tokens 2048, batch_size 240):
5528 tps (original performance) -> 6378 tps (w/ internal bucket size 128)
Add
--bucket_size=128 --bucket_internalto the commands to enable the feature.