huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-08-01 12:50:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Parent links

Browse Source
This commit is contained in:
Daniël de Kok 2024-08-02 13:54:38 +00:00
parent 590fc2c58d
commit 9caee1f368
1 changed files with 13 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
backends/v3/src/radix.rs
View File

@ -24,7 +24,7 @@ pub struct RadixTrie {
impl RadixTrie { impl RadixTrie {
pub fn new() -> Self { pub fn new() -> Self {
let root = TrieNode::new(vec![], vec![], 0); let root = TrieNode::new(vec![], vec![], 0, None);
let mut nodes = SlotMap::new(); let mut nodes = SlotMap::new();
let root = nodes.insert(root); let root = nodes.insert(root);
RadixTrie { RadixTrie {
@ -90,11 +90,11 @@ impl RadixTrie {
let blocks = &blocks[shared_prefix_len..]; let blocks = &blocks[shared_prefix_len..];
self.insert_(child_id, key, blocks) self.insert_(child_id, key, blocks)
} else { } else {
let child = TrieNode::new(key.to_vec(), blocks.to_vec(), self.time); let child = TrieNode::new(key.to_vec(), blocks.to_vec(), self.time, Some(node_id));
let child_id = self.nodes.insert(child); let child_id = self.nodes.insert(child);
let node = self.nodes.get_mut(node_id).unwrap(); let node = self.nodes.get_mut(node_id).unwrap();
node.children.insert(key[0], child_id); node.children.insert(key[0], child_id);
return key.len(); key.len()
} }
} }
@ -105,9 +105,12 @@ impl RadixTrie {
let rest_blocks = node.blocks.split_off(prefix_len); let rest_blocks = node.blocks.split_off(prefix_len);
let first = rest_key[0]; let first = rest_key[0];
let new_id = self let new_id = self.nodes.insert(TrieNode::new(
.nodes rest_key,
.insert(TrieNode::new(rest_key, rest_blocks, self.time)); rest_blocks,
self.time,
Some(node_id),
));
let node = self.nodes.get_mut(node_id).unwrap(); let node = self.nodes.get_mut(node_id).unwrap();
node.children.insert(first, new_id); node.children.insert(first, new_id);
@ -115,19 +118,21 @@ impl RadixTrie {
} }
struct TrieNode { struct TrieNode {
blocks: Vec<u32>,
children: HashMap<u32, NodeId>, children: HashMap<u32, NodeId>,
key: Vec<u32>, key: Vec<u32>,
blocks: Vec<u32>,
last_accessed: u64, last_accessed: u64,
parent: Option<NodeId>,
} }
impl TrieNode { impl TrieNode {
fn new(key: Vec<u32>, blocks: Vec<u32>, last_accessed: u64) -> Self { fn new(key: Vec<u32>, blocks: Vec<u32>, last_accessed: u64, parent: Option<NodeId>) -> Self {
TrieNode { TrieNode {
children: HashMap::new(), children: HashMap::new(),
key, key,
blocks, blocks,
last_accessed, last_accessed,
parent,
} }
} }
} }