Skip to main content

dynamo_kv_router/
lib.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! KV Router - Radix tree data structures for LLM KV cache routing.
5//!
6//! This crate provides the core radix tree implementation and protocols for
7//! efficient KV cache lookup and routing in distributed LLM inference systems.
8
9mod active_set;
10pub(crate) mod cleanup;
11mod lookup_update;
12
13pub mod indexer;
14pub mod protocols;
15pub mod recovery;
16pub mod scheduling;
17pub mod sequences;
18pub mod services;
19pub mod zmq_wire;
20
21// Backward-compat re-exports: old top-level module paths still work
22pub use indexer::concurrent_radix_tree;
23pub use indexer::concurrent_radix_tree_compressed;
24pub use indexer::positional as nested_map;
25pub use indexer::pruning as approx;
26pub use indexer::radix_tree;
27
28pub use scheduling::config;
29pub use scheduling::queue;
30pub use scheduling::selector;
31pub use sequences::multi_worker as multi_worker_sequence;
32pub use sequences::single as sequence;
33
34#[cfg(any(test, feature = "bench"))]
35pub mod test_utils;
36
37// Re-export key types for convenience
38pub use self::multi_worker_sequence::{
39    ActiveSequencesMultiWorker, ReplicaWorkerPolicy, SequenceError, SequencePublisher,
40    SequenceRequest, SequenceSubscriber,
41};
42pub use self::sequence::{ActiveSequences, RequestId};
43pub use self::sequences::{PrefillTokenDeltas, WorkerLoadProjection};
44pub use concurrent_radix_tree::ConcurrentRadixTree;
45pub use concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed;
46pub use config::{
47    KvRouterConfig, RouterConfigOverride, RouterPrefillLoadModel, RouterQueuePolicy,
48    SharedCacheType,
49};
50#[allow(deprecated)]
51pub use indexer::{
52    AnchorAwareBranchShardedIndexer, AnchorRef, AnchorTask, BranchShardedIndexer,
53    LowerTierContinuation, LowerTierIndexer, MaybeError, SharedKvCache, SyncIndexer,
54    ThreadPoolIndexer,
55};
56pub use nested_map::PositionalIndexer;
57pub use protocols::{
58    KvCacheEventError, KvTransferEnforcement, LocalBlockHash, OverlapScores, RouterEvent,
59    RouterEventSink, SharedCacheHits, WorkerConfigLike, WorkerId, compute_block_hash_for_seq,
60};
61pub use queue::SchedulerQueue;
62pub use radix_tree::RadixTree;
63pub use scheduling::LocalScheduler;
64pub use scheduling::PrefillLoadEstimator;
65pub use scheduling::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy};
66pub use scheduling::{KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse};
67pub use selector::{DefaultWorkerSelector, WorkerSelector};