Class RadixCache

java.lang.Object
smile.llm.cache.RadixCache
public class RadixCache extends Object
Radix tree-based KV cache for LLM inference serving.

Organizes KV cache activations in a compressed prefix tree (radix tree) so that requests sharing a common token prefix — such as a shared system prompt — can reuse already-computed KV entries without re-computation.

Each edge in the tree stores a contiguous span of token IDs together with the corresponding GPU KV cache slot indices (held as an int64 Tensor). When a new request arrives the longest matching prefix is found in O(|key|) time, and only the remaining suffix needs to be computed.

Page alignment

When pageSize > 1 all key lengths are rounded down to the nearest multiple of pageSize before matching or inserting, aligning them to hardware KV cache block boundaries.

Namespace isolation

An optional extraKey string (e.g., a LoRA adapter ID or cache salt) can be attached to each insert/match call. Entries with different extraKey values are kept in disjoint subtrees and never share prefix nodes.

Reference counting and eviction

Every node carries a lock reference count. Calling incLockRef(RadixTreeNode) prevents a node and all its ancestors from being evicted; decLockRef(RadixTreeNode) releases the protection. Eviction follows an LRU policy applied to evictable leaves — nodes whose lockRef is zero and that have no live (non-evicted) children.

Tensor ownership

Each RadixTreeNode owns its RadixTreeNode.value tensor and is responsible for closing it when the node is evicted or split. The MatchResult returned by matchPrefix(int[]) is AutoCloseable: use try-with-resources to release the underlying tensor automatically.
See Also:

  • Constructor Summary

    Constructors
    Constructor
    Description
    RadixCache()
    Constructs a cache with the default page size of 1.
    RadixCache(int pageSize)
    Constructor.

  • Method Summary

    Modifier and Type
    Method
    Description
    int
    decLockRef(RadixTreeNode node)
    Decrements the lock reference count on node and all its ancestors up to (but not including) the root.
    int
    evict(int numTokens, Consumer<Tensor> freeCallback)
    Evicts LRU leaf nodes until at least numTokens KV cache token slots have been freed, calling freeCallback for each evicted node's Tensor.
    int
    evictableSize()
    Returns the total number of evictable (unlocked) tokens currently cached.
    int
    incLockRef(RadixTreeNode node)
    Increments the lock reference count on node and all its ancestors up to (but not including) the root, protecting them from eviction.
    InsertResult
    insert(int[] tokenIds, Tensor kvIndices)
    Inserts a token sequence and its KV cache slot indices into the radix tree.
    InsertResult
    insert(int[] tokenIds, Tensor kvIndices, String extraKey, int priority)
    Inserts a token sequence and its KV cache slot indices into the radix tree.
    MatchResult
    matchPrefix(int[] tokenIds)
    Finds the longest cached prefix of the given token sequence.
    MatchResult
    matchPrefix(int[] tokenIds, String extraKey)
    Finds the longest cached prefix of the given token sequence within a namespace.
    void
    prettyPrint()
    Prints the radix tree structure to stdout for debugging.
    int
    protectedSize()
    Returns the total number of protected (locked) tokens currently cached.
    void
    reset()
    Resets the cache, discarding all stored KV data and reference counts.
    int
    totalSize()
    Returns the total number of tokens in the cache (evictable + protected).

    Methods inherited from class Object

    clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

  • Constructor Details

    • RadixCache

      public RadixCache(int pageSize)
      Constructor.
      Parameters:
      pageSize - the number of tokens per KV cache page (must be ≥ 1).
      Throws:
      IllegalArgumentException - if pageSize < 1.

    • RadixCache

      public RadixCache()
      Constructs a cache with the default page size of 1.

  • Method Details

    • reset

      public void reset()
      Resets the cache, discarding all stored KV data and reference counts. Every live RadixTreeNode.value tensor is closed before the tree is cleared. The root sentinel is re-created with an empty key.

    • matchPrefix

      public MatchResult matchPrefix(int[] tokenIds)
      Finds the longest cached prefix of the given token sequence. Equivalent to calling matchPrefix(int[], String) with a null extra key.
      Parameters:
      tokenIds - the token ID sequence to look up.
      Returns:
      the match result. Use try-with-resources to release the enclosed tensor automatically, or call MatchResult.close() explicitly.

    • matchPrefix

      public MatchResult matchPrefix(int[] tokenIds, String extraKey)
      Finds the longest cached prefix of the given token sequence within a namespace.

      The search is limited to nodes whose extraKey matches the supplied value, so entries from different namespaces are never confused.

      If the match ends inside a stored node (partial match), the node is automatically split at the boundary, improving future lookups. The evictable size is unchanged because the total number of tokens is conserved.

      This call refreshes the RadixTreeNode.lastAccessTime of every visited node, feeding the LRU eviction policy.

      Parameters:
      tokenIds - the token ID sequence to look up.
      extraKey - optional namespace tag, or null for the default namespace.
      Returns:
      the match result. Use try-with-resources to release the enclosed tensor automatically, or call MatchResult.close() explicitly.

    • insert

      public InsertResult insert(int[] tokenIds, Tensor kvIndices)
      Inserts a token sequence and its KV cache slot indices into the radix tree. Equivalent to calling insert(int[], Tensor, String, int) with a null extra key and priority 0.
      Parameters:
      tokenIds - the token ID sequence.
      kvIndices - a 1-D int64 tensor of KV cache slot indices, one per token (kvIndices.length() >= tokenIds.length).
      Returns:
      the insert result.

    • insert

      public InsertResult insert(int[] tokenIds, Tensor kvIndices, String extraKey, int priority)
      Inserts a token sequence and its KV cache slot indices into the radix tree.

      The insert walks the tree consuming as many tokens as match existing edges. Existing nodes that partially overlap with the new key are split. The page-aligned portion of the key not yet in the tree is stored as a new leaf node backed by a freshly-allocated int64 tensor.

      The returned InsertResult.prefixLen() indicates how many leading tokens were already cached. The caller can free the duplicate KV slot entries for those positions (indices [0, prefixLen)).

      Parameters:
      tokenIds - the token ID sequence to cache.
      kvIndices - a 1-D int64 tensor of KV cache slot indices (kvIndices.length() >= tokenIds.length).
      extraKey - optional namespace tag, or null for the default namespace.
      priority - eviction priority; higher values delay eviction.
      Returns:
      the insert result.

    • evict

      public int evict(int numTokens, Consumer<Tensor> freeCallback)
      Evicts LRU leaf nodes until at least numTokens KV cache token slots have been freed, calling freeCallback for each evicted node's Tensor.

      Eviction always starts from the least-recently-used evictable leaf. When a leaf is evicted its parent may itself become a new evictable leaf (if all other children were already evicted) and is then eligible for further eviction in the same call.

      The freeCallback receives ownership of the tensor and is responsible for releasing it (e.g., by calling Tensor.close()). If freeCallback is null the tensor is closed immediately.

      Parameters:
      numTokens - the minimum number of token slots to reclaim.
      freeCallback - invoked with the int64 KV index tensor of each evicted node. May be null.
      Returns:
      the actual number of token slots freed (may exceed numTokens).

    • incLockRef

      public int incLockRef(RadixTreeNode node)
      Increments the lock reference count on node and all its ancestors up to (but not including) the root, protecting them from eviction.

      Nodes transition from the evictable pool to the protected pool as their lock count becomes positive for the first time.

      Parameters:
      node - the node to protect (may be null, which is a no-op).
      Returns:
      the change in evictableSize (always ≤ 0).

    • decLockRef

      public int decLockRef(RadixTreeNode node)
      Decrements the lock reference count on node and all its ancestors up to (but not including) the root. When a node's count reaches zero it re-enters the evictable pool and becomes eligible for LRU eviction.
      Parameters:
      node - the node to release (may be null, which is a no-op).
      Returns:
      the change in evictableSize (always ≥ 0).

    • evictableSize

      public int evictableSize()
      Returns the total number of evictable (unlocked) tokens currently cached.
      Returns:
      evictable token count.

    • protectedSize

      public int protectedSize()
      Returns the total number of protected (locked) tokens currently cached.
      Returns:
      protected token count.

    • totalSize

      public int totalSize()
      Returns the total number of tokens in the cache (evictable + protected). Traverses the live (non-evicted) subtree iteratively.
      Returns:
      total cached token count.

    • prettyPrint

      public void prettyPrint()
      Prints the radix tree structure to stdout for debugging.