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.

Sign up for GitHub

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

Jump to bottom

TriePagedAttentionCache #632

Merged
merged 23 commits into from
Dec 4, 2024
Merged

TriePagedAttentionCache #632

merged 23 commits into from
Dec 4, 2024
+943 −40

Conversation

renxida
Copy link
Contributor

@renxida renxida commented Dec 2, 2024 •
edited
Loading

feat: Add TriePagedAttentionCache with initial implementation

Added TriePagedAttentionCache as an optional prefix sharing algorithm, selectable via:
config["paged_kv_cache"]["prefix_sharing_algorithm"] = "trie"

Current Status:

  • Basic implementation and unit tests complete
  • Integration test cases for both Base and Trie implementations, with trie implementation xfailed due to pending cache allocation improvements
  • BasePagedAttentionCache remains the default

Next Steps:
To achieve full functionality, we need to support cache re-allocations to extend the associated tokens & pages.

@renxida renxida marked this pull request as ready for review December 2, 2024 22:24
@renxida renxida requested a review from stbaione December 2, 2024 22:24
@renxida
Copy link
Contributor Author

renxida commented Dec 2, 2024

image

: D

@renxida renxida mentioned this pull request Dec 3, 2024
Open
20 tasks
stbaione
stbaione approved these changes Dec 3, 2024
View reviewed changes
Copy link
Contributor

@stbaione stbaione 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 this could benefit from another set of eyes, but the overall interface and operations make sense to me

@renxida
add unit tests and fix numerous small problems
a1946b8
@renxida
all tests passing
8437a7f
@renxida
all but publishing working
24a4313
@renxida
naming consistency Page -> Paged
5c87b2b
@renxida
passing all tests now
07d3040
@renxida
name and documentation update publish_pages -> publish_pages_for_tokens
8f83118
@renxida
undo accidental edit of messages.py
5c37666
@renxida
remove not-very-useful test case
802c4e2
@renxida
fix missing rename
3c7cc6c
@renxida
use trie by default
f59485e
@renxida
add errmsg for unknown kvcache type
646c167
@renxida
fix typoes
5801d27
@renxida
trie doesn't work yet reverting to base for default
7f6e5a3
@renxida
rename cache_type to prefix_sharing_algorithm
ca29397
@renxida
add another test case for trie and xfail it
7abaa7c
@renxida
fix xpass by fixing a place where i forgot to pass prefix_sharing_alg…
1e35634 
…orithm in conftest.py
@renxida
organize test cases
1a419f1
@renxida
add an extend_allocation method to cache allocations and make decode …
85bc5fc 
…requests use that instead
@renxida
remove xfail because we are now PASSING!
feb3aa9
@renxida
fix typo; force certain args to be kwargs
7e07cc2
@renxida
use a reference counter class to replace int for ref_count in trie node
d71826f
@renxida renxida enabled auto-merge (squash) December 4, 2024 04:56
@renxida renxida merged commit de4d2fe into nod-ai:main Dec 4, 2024
13 of 20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@stbaione stbaione stbaione approved these changes

@rsuderman rsuderman rsuderman approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@renxida @rsuderman @stbaione