lib.rs 1.84 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;
Yan Ru Pei's avatar
Yan Ru Pei committed
10
pub mod concurrent_radix_tree;
11
pub mod event_sink;
12
pub mod indexer;
13
14
#[cfg(feature = "bench")]
pub mod naive_indexers;
Yan Ru Pei's avatar
Yan Ru Pei committed
15
pub mod nested_map;
16
pub mod protocols;
17
pub mod radix_tree;
18
19
pub mod scheduling;
pub mod sequences;
20
pub mod zmq_wire;
21
22
23
24
25
26
27

// Backward-compat re-exports: preserve old module paths for external consumers
pub use scheduling::config;
pub use scheduling::queue;
pub use scheduling::selector;
pub use sequences::multi_worker as multi_worker_sequence;
pub use sequences::single as sequence;
28

29
30
#[cfg(any(test, feature = "bench"))]
pub mod test_utils;
Yan Ru Pei's avatar
Yan Ru Pei committed
31

32
// Re-export key types for convenience
33
pub use self::multi_worker_sequence::{
34
35
36
    ActiveSequencesMultiWorker, SequenceError, SequencePublisher, SequenceRequest,
    SequenceSubscriber,
};
37
38
39
pub use self::sequence::{ActiveSequences, RequestId};
pub use concurrent_radix_tree::ConcurrentRadixTree;
pub use config::{KvRouterConfig, RouterConfigOverride};
40
pub use event_sink::EventSink;
41
pub use indexer::{MaybeError, SyncIndexer, ThreadPoolIndexer};
42
43
#[cfg(feature = "bench")]
pub use naive_indexers::{InvertedIndex, NaiveNestedMap};
Yan Ru Pei's avatar
Yan Ru Pei committed
44
pub use nested_map::PositionalIndexer;
45
pub use protocols::{
46
    KvCacheEventError, LocalBlockHash, OverlapScores, RouterEvent, WorkerConfigLike, WorkerId,
47
48
    compute_block_hash_for_seq,
};
49
pub use queue::SchedulerQueue;
50
pub use radix_tree::RadixTree;
51
52
pub use scheduling::{KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse};
pub use selector::{DefaultWorkerSelector, WorkerSelector};