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
9use std::sync::Arc;
10
11mod active_set;
12pub(crate) mod cleanup;
13pub mod conditional_disagg;
14mod lookup_update;
15
16pub mod identity;
17pub mod indexer;
18pub mod protocols;
19pub mod recovery;
20pub mod router_hint;
21pub mod scheduling;
22pub mod sequences;
23pub mod services;
24pub mod tracking_hash;
25pub mod worker_type;
26pub mod zmq_wire;
27
28// Backward-compat re-exports: old top-level module paths still work
29pub use indexer::concurrent_radix_tree;
30pub use indexer::concurrent_radix_tree_compressed;
31pub use indexer::positional as nested_map;
32pub use indexer::pruning as approx;
33pub use indexer::radix_tree;
34
35pub use scheduling::config;
36pub use scheduling::queue;
37pub use scheduling::selector;
38pub use sequences::multi_worker as multi_worker_sequence;
39pub use sequences::single as sequence;
40
41#[cfg(any(test, feature = "bench"))]
42pub mod test_utils;
43
44// Re-export key types for convenience
45pub use self::multi_worker_sequence::{
46    ActiveSequencesMultiWorker, NoopSequencePublisher, ReplicaWorkerPolicy, SequenceError,
47    SequencePublisher, SequenceRequest, SequenceSubscriber,
48};
49pub use self::sequence::{ActiveSequences, RequestId};
50pub use self::sequences::{PrefillTokenDeltas, WorkerLoadProjection};
51pub use concurrent_radix_tree::ConcurrentRadixTree;
52pub use concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed;
53pub use config::{
54    ConditionalDisaggPolicyKind, KvRouterConfig, RouterConfigOverride, RouterPrefillLoadModel,
55    RouterQueuePolicy, SharedCacheType,
56};
57pub use identity::{DEFAULT_ROUTING_GROUP, DcId, RoutingPartitionId, RoutingPartitionRef};
58#[allow(deprecated)]
59pub use indexer::{
60    AnchorAwareBranchShardedIndexer, AnchorRef, AnchorTask, BranchShardedIndexer,
61    LowerTierContinuation, LowerTierIndexer, MaybeError, SharedKvCache, SyncIndexer,
62    ThreadPoolIndexer,
63};
64pub use nested_map::PositionalIndexer;
65pub use protocols::{
66    KvCacheEventError, KvTransferEnforcement, LocalBlockHash, OverlapScores, RouterEvent,
67    RouterEventSink, SharedCacheHits, WorkerConfigLike, WorkerId, compute_block_hash_for_seq,
68};
69pub use queue::SchedulerQueue;
70pub use radix_tree::RadixTree;
71pub use scheduling::LocalScheduler;
72pub use scheduling::PrefillLoadEstimator;
73pub use scheduling::policy::{FcfsPolicy, RouterSchedulingPolicy, SchedulingPolicy, WsptPolicy};
74pub use scheduling::{
75    KvSchedulerError, PotentialLoad, SchedulingRequest, SchedulingResponse, SessionContext,
76    WorkerSelectionInputTrigger, WorkerSelectionKvHints, WorkerSelectionPolicyError,
77};
78pub use selector::{
79    DefaultWorkerSelector, ScoredWorkerCandidate, WorkerCacheInput, WorkerCandidate, WorkerFilter,
80    WorkerInputView, WorkerInputs, WorkerLoadInput, WorkerPicker, WorkerScorer,
81    WorkerSelectionContext, WorkerSelectionInput, WorkerSelectionPolicy, WorkerSelector,
82};
83pub use tracking_hash::{TrackingHashAlgorithm, TrackingHashContext, TrackingHashScope};
84pub use worker_type::WorkerType;
85
86/// Factory that creates one worker-selection policy per routing partition.
87pub type WorkerSelectionPolicyFactory = Arc<
88    dyn for<'a> Fn(&KvRouterConfig, WorkerType, RoutingPartitionRef<'a>) -> WorkerSelectionPolicy
89        + Send
90        + Sync,
91>;
dynamo_kv_router/lib.rs

dynamo_kv_router/
lib.rs
