enable internal kv bucket in llama

[xt574chen]

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_internal to the commands to enable the feature.

[HuggingFaceDocBuilderDev]

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.

[regisss]

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

[xt574chen]

@puneeshkhanna I see you've corrected some error messages. I hope this update won't cause conflict.

[puneeshkhanna]

@xt574chen - I will update my PR once this gets merged first

[regisss]

@xt574chen Can you run the following from the root of the repo to make the code style check pass please?

pip install -U ruff
make style