huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid continuous sorting during reclamation

Browse Source
This commit is contained in:
Daniël de Kok 2024-07-10 20:12:14 +02:00
parent c352a3e231
commit 1a461234d5
1 changed files with 29 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
router/src/infer/v3/block_allocator.rs
View File

@ -1,8 +1,8 @@
use itertools::Itertools; use itertools::Itertools;
use std::{ use std::{
borrow::BorrowMut, borrow::BorrowMut,
cmp::min, cmp::{min, Reverse},
collections::{hash_map::Entry, HashMap, HashSet}, collections::{hash_map::Entry, BinaryHeap, HashMap, HashSet},
hash::{DefaultHasher, Hash, Hasher}, hash::{DefaultHasher, Hash, Hasher},
}; };
use tokio::sync::{mpsc, oneshot}; use tokio::sync::{mpsc, oneshot};
@ -257,7 +257,7 @@ impl PrefixCache {
}) })
} }
fn free_prefix(&mut self, prefix_hash: u64) { fn free_prefix_block(&mut self, prefix_hash: u64) {
let state = self let state = self
.cache_blocks .cache_blocks
.remove(&prefix_hash) .remove(&prefix_hash)
@ -269,6 +269,7 @@ impl PrefixCache {
} }
self.leaves.remove(&prefix_hash); self.leaves.remove(&prefix_hash);
self.free_blocks.push(state.block_id);
} }
fn decref_prefix(&mut self, prefix_hash: u64) { fn decref_prefix(&mut self, prefix_hash: u64) {
@ -293,36 +294,40 @@ impl PrefixCache {
} }
fn alloc_or_reclaim(&mut self, n_tokens: usize) -> Option<Vec<u32>> { fn alloc_or_reclaim(&mut self, n_tokens: usize) -> Option<Vec<u32>> {
let n_blocks = (n_tokens + self.block_size - 1) / self.block_size; let n_blocks_needed = (n_tokens + self.block_size - 1) / self.block_size;
let n_blocks_needed = if n_blocks > self.free_blocks.len() {
n_blocks - self.free_blocks.len()
} else {
0
};
while self.free_blocks.len() < n_blocks_needed { let mut reclaimable_blocks = self
// We have to free one block at a time, because removing the LRU
// prefix block may make available another prefix block that is
// LRU.
//
// TODO: switch to something like a binary heap to avoid sorting
// the set of leaves over and over again.
let (lru_prefix_hash, lru_block_id) = self
.leaves .leaves
.iter() .iter()
.map(|prefix_hash| (prefix_hash, &self.cache_blocks[prefix_hash])) .map(|prefix_hash| {
.sorted_by_key(|state| state.1.last_accessed) let state = &self.cache_blocks[prefix_hash];
.map(|(prefix_hash, state)| (*prefix_hash, state.block_id)) Reverse((state.last_accessed, *prefix_hash, state.predecessor))
.next()?; })
.collect::<BinaryHeap<_>>();
self.free_prefix(lru_prefix_hash); while self.free_blocks.len() < n_blocks_needed {
self.free_blocks.push(lru_block_id); // We have to free one block at a time because removing the LRU
// prefix block may make available another prefix block that is
// LRU.
let (_, lru_prefix_hash, predecessor) = reclaimable_blocks.pop()?.0;
self.free_prefix_block(lru_prefix_hash);
// TODO: this is a leaky abstraction, avoid this.
if let Some(predecessor) = predecessor {
let state = &self.cache_blocks[&predecessor];
if state.ref_count == 0 {
reclaimable_blocks.push(Reverse((
state.last_accessed,
predecessor,
state.predecessor,
)));
}
}
} }
Some( Some(
self.free_blocks self.free_blocks
.split_off(self.free_blocks.len() - n_blocks), .split_off(self.free_blocks.len() - n_blocks_needed),
) )
} }