This PR adds the feature requested in #30758. The HHCache class is almost directly taken from the original H2O paper's authors code found here. Currently the PR only adds the changes required to Llama model class. As of now I have taken @gante 's suggestion of adding Cache.post_process() and calling it within LlamaAttention.forward.
To-Do
I'm not sure if the logic for RoPE rerotation is 100% correct. I think the recent tokens are correct, but not the hh tokens after eviction. Would love to have another set of eyes on that.
Write tests to ensure that this HHCache class has the same behavior compared to the original code by paper authors.
Benchmarking(?)
Feedback and/or help would be appreciated. Thanks!
amyeroberts
commented
2 days ago
What does this PR do?
This PR adds the feature requested in #30758. The HHCache class is almost directly taken from the original H2O paper's authors code found here. Currently the PR only adds the changes required to Llama model class. As of now I have taken @gante 's suggestion of adding
Cache.post_process()and calling it withinLlamaAttention.forward.To-Do
Feedback and/or help would be appreciated. Thanks!