Improve turbomind's prefix cache #3332
Open
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
|
性能提高多少？ |
lzhangzz
reviewed
Jun 26, 2025
| @@ -1011,16 +958,16 @@ void LlamaBatch::OutputLogits(const Tensor& logits, int first, int last, Generat | |||
|
|
|||
| int diff = (history_len + offset) - cache_len; | |||
|
|
|||
| const int valid_len = input_len - std::max(0, (history_len + offset) - cache_len); | |||
| const int valid_len = input_len - std::max(0, diff); | |||
There was a problem hiding this comment.
do not replace (history_len + offset) - cache_len with diff, it's harder to understand the code. diff is only used to print debug info.
| freed_.insert(freed_.end(), seq.blocks.begin(), seq.blocks.end()); | ||
| } | ||
| it = sequences_.erase(it); | ||
| else { |
There was a problem hiding this comment.
The else branch has nothing to do with the function parameter and is going to be called multiple times when multiple sequence are erased.
| BlockIds block_ids; | ||
| UniqueIds block_unique_ids; | ||
| std::vector<std::shared_ptr<TrieNode>> nodes; | ||
| std::tie(block_ids, block_unique_ids, nodes) = block_trie_->Cache(seq, seq.prompt); |
There was a problem hiding this comment.
Consider prune the tree while matching prefixes instead of the match-verify-remove pattern
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.
This Pull Request implements several improvements to the tm engine:
enable_prefix_cachingdefaults to TrueRemaining refactoring will be processed in another PR:
get_promptin model.pyapply_chat_templateinsteadtest cases:
https://aicarrier.feishu.cn/wiki/JxoDwiOO7i0xnvkZHzxc3PdGnJh?sheet=95YWnW