Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding cached KVs #266

Merged
merged 16 commits into from
May 22, 2023
Merged

Adding cached KVs #266

merged 16 commits into from
May 22, 2023

Conversation

gkroiz
Copy link
Contributor

@gkroiz gkroiz commented May 13, 2023

This PR adds the option to cache KVs for inference.

Cached KVs portion of splitting #240 into 2 PRs

main changes:

added KV caching
moved rope cache to the whole model from each attention layer
added mask cache to use for attention masking with cached kvs.
fix ops and tensors that would cause recompilation.

Similar to Lightning-AI/litgpt#51
Fixes #197
Fixes #309

TODO: Update tests
@gkroiz
Copy link
Contributor Author

gkroiz commented May 13, 2023

TODO: update tests

@lantiga
Copy link
Collaborator

lantiga commented May 16, 2023

For visibility: updates to lit-parrot caching
Lightning-AI/litgpt#60

Copy link
Contributor

@carmocca carmocca left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Implementation looks good!

@gkroiz
Copy link
Contributor Author

gkroiz commented May 22, 2023

@carmocca what is the difference between generate.py and generate/full.py?

@carmocca
Copy link
Contributor

@gkroiz It's basically a duplicate. See the description here #255

@carmocca
Copy link
Contributor

I opened #313 to track it. I wouldn't try to merge them here.

Copy link
Contributor

@carmocca carmocca left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we're done here! The howto update can be done separately

@carmocca carmocca self-assigned this May 22, 2023
@carmocca carmocca added the enhancement New feature or request label May 22, 2023
Copy link
Collaborator

@lantiga lantiga left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome job!

@lantiga lantiga merged commit a24fc5e into Lightning-AI:main May 22, 2023
@gkroiz gkroiz deleted the cached_kvs branch May 23, 2023 02:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request inference
Projects
None yet
Development

Successfully merging this pull request may close these issues.

when will kv caching landing? KV cache for faster generation
3 participants