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;
11pub mod conditional_disagg;
12mod lookup_update;
13
14pub mod identity;
15pub mod indexer;
16pub mod protocols;
17pub mod recovery;
18pub mod scheduling;
19pub mod sequences;
20pub mod services;
21pub mod tracking_hash;
22pub mod zmq_wire;
23
24// Backward-compat re-exports: old top-level module paths still work
25pub use indexer::concurrent_radix_tree;
26pub use indexer::concurrent_radix_tree_compressed;
27pub use indexer::positional as nested_map;
28pub use indexer::pruning as approx;
29pub use indexer::radix_tree;
30
31pub use scheduling::config;
32pub use scheduling::queue;
33pub use scheduling::selector;
34pub use sequences::multi_worker as multi_worker_sequence;
35pub use sequences::single as sequence;
36
37#[cfg(any(test, feature = "bench"))]
38pub mod test_utils;
39
40// Re-export key types for convenience
41pub use self::multi_worker_sequence::{
42    ActiveSequencesMultiWorker, NoopSequencePublisher, ReplicaWorkerPolicy, SequenceError,
43    SequencePublisher, SequenceRequest, SequenceSubscriber,
44};
45pub use self::sequence::{ActiveSequences, RequestId};
46pub use self::sequences::{PrefillTokenDeltas, WorkerLoadProjection};
47pub use concurrent_radix_tree::ConcurrentRadixTree;
48pub use concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed;
49pub use config::{
50    ConditionalDisaggPolicyKind, KvRouterConfig, RouterConfigOverride, RouterPrefillLoadModel,
51    RouterQueuePolicy, SharedCacheType,
52};
53pub use identity::{DEFAULT_ROUTING_GROUP, DcId, RoutingPartitionId, RoutingPartitionRef};
54#[allow(deprecated)]
55pub use indexer::{
56    AnchorAwareBranchShardedIndexer, AnchorRef, AnchorTask, BranchShardedIndexer,
57    LowerTierContinuation, LowerTierIndexer, MaybeError, SharedKvCache, SyncIndexer,
58    ThreadPoolIndexer,
59};
60pub use nested_map::PositionalIndexer;
61pub use protocols::{
62    KvCacheEventError, KvTransferEnforcement, LocalBlockHash, OverlapScores, RouterEvent,
63    RouterEventSink, SharedCacheHits, WorkerConfigLike, WorkerId, compute_block_hash_for_seq,
64};
65pub use queue::SchedulerQueue;
66pub use radix_tree::RadixTree;
67pub use scheduling::LocalScheduler;
68pub use scheduling::PrefillLoadEstimator;
69pub use scheduling::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy};
70pub use scheduling::{KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse};
71pub use selector::{DefaultWorkerSelector, WorkerSelector};
72pub use tracking_hash::{TrackingHashAlgorithm, TrackingHashContext, TrackingHashScope};