Skip to main content

/

Docs

Nightly

v24.6

Log in

Python module

radix_trie

RadixTrie​

class max.pipelines.kv_cache.radix_trie.RadixTrie

insert()​

insert(tokens: List[Any], blocks: List[Any], node: TrieNode | None = None) → TrieNode

Inserts tokens and blocks into the trie.

We assume that each block contains exactly one token so the length of both input lists must match.

  • Parameters:

    • tokens – Tokens to insert into trie
    • blocks – KV cache block for each token
    • node – Node to begin insertion at. By default, insertion begins at root node

  • Returns:

    Node corresponding to end of the sequence where future : generated tokens can be inserted

  • Return type:

    trie_node

match_prefix()​

match_prefix(tokens: List[Any]) → Tuple[TrieNode, List[Any]]

Matches the input tokens with the contents of the trie.

  • Parameters:

    tokens – tokens to search the trie for

  • Returns:

    • trie_node: Node corresponding to end of matched prefix where : future generated tokens can be inserted
    • block_list: KV cache blocks for matched prefix

  • Return type:

    Tuple containing

pretty_format()​

pretty_format() → List[str]

Formats the contents of the trie.

TrieNode​

class max.pipelines.kv_cache.radix_trie.TrieNode