implement of DISCO #10904
Open
implement of DISCO #10904
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…LaMA - Add DISCO configuration parameters to LlamaConfig - Create DISCOCache class for intelligent token eviction - Integrate DISCO into LlamaAttention and LlamaModel - Add comprehensive unit tests and examples - Include detailed documentation DISCO reduces KV cache memory to ~3.2% while maintaining model quality through adaptive layer-wise allocation and attention-based scoring.
|
Thanks for your contribution! |
- Fixed index_select error where indices exceeded sequence length - Changed eviction logic to properly handle scores in sequence dimension - Average scores across batch and heads before selecting top-k tokens - Ensure indices are within valid range for the sequence The issue was that we were flattening all dimensions and getting indices that could exceed the sequence length. Now we properly average across batch and head dimensions first, then select top-k tokens from the sequence dimension only.
- Fixed eviction logic to not evict on first update (initialization) - Corrected score averaging across batch and head dimensions - Ensured indices are within valid sequence range - Added standalone test script to verify all DISCO functionality The key changes: 1. First cache update now preserves all tokens (no eviction on init) 2. Eviction only happens on subsequent updates when budget is exceeded 3. Scores are properly averaged before selecting top-k tokens 4. All 5 DISCO tests now pass successfully
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Before submitting
testsfolder. If there are codecov issues, please add tests cases first.PR types
PR changes
Description
KV Cache压缩算法DISCO算法的实现。