lib.rs 1.29 KB
Newer Older
1
2
3
4
5
6
7
8
9
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! KV Router - Radix tree data structures for LLM KV cache routing.
//!
//! This crate provides the core radix tree implementation and protocols for
//! efficient KV cache lookup and routing in distributed LLM inference systems.

pub mod approx;
10
11
#[cfg(feature = "bench")]
pub mod bench_utils;
Yan Ru Pei's avatar
Yan Ru Pei committed
12
pub mod concurrent_radix_tree;
13
pub mod indexer;
14
pub mod multi_worker_sequence;
15
16
#[cfg(feature = "bench")]
pub mod naive_indexers;
Yan Ru Pei's avatar
Yan Ru Pei committed
17
pub mod nested_map;
18
pub mod protocols;
19
pub mod radix_tree;
20
pub mod sequence;
21

Yan Ru Pei's avatar
Yan Ru Pei committed
22
23
24
#[cfg(test)]
pub(crate) mod test_utils;

25
// Re-export key types for convenience
Yan Ru Pei's avatar
Yan Ru Pei committed
26
27
pub use concurrent_radix_tree::ConcurrentRadixTree;
pub use indexer::{MaybeError, SyncIndexer, ThreadPoolIndexer};
28
29
30
31
pub use multi_worker_sequence::{
    ActiveSequencesMultiWorker, SequenceError, SequencePublisher, SequenceRequest,
    SequenceSubscriber,
};
32
33
#[cfg(feature = "bench")]
pub use naive_indexers::{InvertedIndex, NaiveNestedMap};
Yan Ru Pei's avatar
Yan Ru Pei committed
34
pub use nested_map::PositionalIndexer;
35
36
37
38
39
pub use protocols::{
    KvCacheEventError, LocalBlockHash, OverlapScores, RouterEvent, WorkerId,
    compute_block_hash_for_seq,
};
pub use radix_tree::RadixTree;
40
pub use sequence::{ActiveSequences, RequestId};