kv_router.rs 39.4 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
use std::collections::HashMap;
5
use std::sync::Arc;
6
use std::time::Duration;
7

8
use anyhow::Result;
9
use derive_builder::Builder;
10
use dynamo_runtime::{
11
    component::{Client, Endpoint},
12
    discovery::{DiscoveryQuery, EventTransportKind},
13
    pipeline::{
14
15
        AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, PushRouter, ResponseStream,
        SingleIn, async_trait,
16
    },
17
    protocols::EndpointId,
18
    protocols::annotated::Annotated,
19
    traits::DistributedRuntimeProvider,
20
21
};
use futures::stream::{self, StreamExt};
22
use rand::Rng;
23
use serde::{Deserialize, Serialize};
24
use serde_json::json;
25

26
27
28
29
30
// Re-export from dynamo-kv-router crate
pub use dynamo_kv_router::approx;
pub use dynamo_kv_router::indexer;
pub use dynamo_kv_router::protocols;

31
pub mod prefill_router;
32
pub mod publisher;
33
pub mod recorder;
34
pub mod scheduler;
35
pub mod sequence;
36
pub mod subscriber;
37
pub mod worker_query;
38

39
use indexer::WorkerKvQueryResponse;
40
pub use prefill_router::PrefillRouter;
41
use worker_query::WorkerQueryClient;
42

43
use crate::{
44
    discovery::RuntimeConfigs,
45
    kv_router::{
46
        approx::PruneConfig,
47
        indexer::{KvIndexer, KvIndexerInterface, KvRouterError},
Yan Ru Pei's avatar
Yan Ru Pei committed
48
        protocols::{
49
            DpRank, LocalBlockHash, OverlapScores, RouterEvent, RouterRequest, RouterResponse,
50
51
            TokensWithHashes, WorkerId, WorkerSelectionResult, WorkerWithDpRank,
            compute_block_hash_for_seq, compute_seq_hash_for_block,
Yan Ru Pei's avatar
Yan Ru Pei committed
52
        },
53
        scheduler::{KvScheduler, KvSchedulerError, PotentialLoad, SchedulingRequest},
54
        sequence::SequenceError,
55
        subscriber::{start_kv_router_background, start_kv_router_background_event_plane},
56
    },
57
    local_model::runtime_config::ModelRuntimeConfig,
58
    preprocessor::PreprocessedRequest,
59
    protocols::common::llm_backend::LLMEngineOutput,
60
    protocols::common::timing::RequestPhase,
61
62
};

63
64
// [gluo TODO] shouldn't need to be public
// this should be discovered from the component
65
66
67
68
69

// for metric scraping (pull-based)
pub const KV_METRICS_ENDPOINT: &str = "load_metrics";

// for metric publishing (push-based)
70
pub const KV_EVENT_SUBJECT: &str = "kv-events";
71
pub const KV_HIT_RATE_SUBJECT: &str = "kv-hit-rate";
72
73
74
75
76
pub const KV_METRICS_SUBJECT: &str = "kv_metrics";

// for inter-router comms
pub const PREFILL_SUBJECT: &str = "prefill_events";
pub const ACTIVE_SEQUENCES_SUBJECT: &str = "active_sequences_events";
77

78
79
80
81
// for radix tree snapshot storage
pub const RADIX_STATE_BUCKET: &str = "radix-bucket";
pub const RADIX_STATE_FILE: &str = "radix-state";

82
83
84
// for worker-local kvindexer query
pub const WORKER_KV_INDEXER_BUFFER_SIZE: usize = 1024; // store 1024 most recent events in worker buffer

85
86
87
88
89
90
/// Generates a dp_rank-specific endpoint name for the worker KV indexer query service.
/// Each dp_rank has its own LocalKvIndexer and query endpoint to ensure per-dp_rank monotonicity.
pub fn worker_kv_indexer_query_endpoint(dp_rank: DpRank) -> String {
    format!("worker_kv_indexer_query_dp{dp_rank}")
}

91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
// for router discovery registration
pub const KV_ROUTER_COMPONENT: &str = "kv-router";
pub const KV_ROUTER_ENDPOINT: &str = "generate";

/// Creates an EndpointId for the KV router in the given namespace.
pub fn router_endpoint_id(namespace: String) -> EndpointId {
    EndpointId {
        namespace,
        component: KV_ROUTER_COMPONENT.to_string(),
        name: KV_ROUTER_ENDPOINT.to_string(),
    }
}

/// Creates a DiscoveryQuery for the KV router in the given namespace.
pub fn router_discovery_query(namespace: String) -> DiscoveryQuery {
    DiscoveryQuery::Endpoint {
        namespace,
        component: KV_ROUTER_COMPONENT.to_string(),
        endpoint: KV_ROUTER_ENDPOINT.to_string(),
    }
}

113
114
115
116
/// A trait that users can implement to define custom selection logic
pub trait WorkerSelector {
    fn select_worker(
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
117
        workers: &HashMap<protocols::WorkerId, Option<ModelRuntimeConfig>>,
118
        request: &SchedulingRequest,
119
        block_size: u32,
120
121
    ) -> Result<WorkerSelectionResult, KvSchedulerError>;
}
122

123
124
125
126
127
128
129
130
131
132
/// Override configuration for router settings that can be specified per-request
#[derive(Debug, Clone, Default, Builder, Serialize, Deserialize)]
pub struct RouterConfigOverride {
    #[builder(default)]
    pub overlap_score_weight: Option<f64>,

    #[builder(default)]
    pub router_temperature: Option<f64>,
}

133
/// KV Router configuration parameters
134
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
135
136
137
pub struct KvRouterConfig {
    pub overlap_score_weight: f64,

138
    pub router_temperature: f64,
139

140
141
    pub use_kv_events: bool,

142
143
    pub router_replica_sync: bool,

144
145
146
    /// Whether to track active blocks in the router (default: true)
    pub router_track_active_blocks: bool,

147
148
149
150
151
    /// Whether to track output blocks during generation (default: false)
    /// When enabled, the router adds placeholder blocks as tokens are generated
    /// and applies fractional decay based on progress toward expected_output_tokens.
    pub router_track_output_blocks: bool,

152
153
154
155
156
    /// Whether to assume KV cache reuse when tracking active blocks (default: true).
    /// When true, computes actual block hashes for sequence tracking.
    /// When false, generates random hashes (assuming no KV cache reuse).
    pub router_assume_kv_reuse: bool,

157
158
159
    /// Threshold for triggering snapshots. If None, no snapshots will be performed.
    pub router_snapshot_threshold: Option<u32>,

160
    /// Whether to reset the router state on startup (default: false)
161
    pub router_reset_states: bool,
162
163
164
165

    /// TTL for blocks in seconds (only used when use_kv_events is false, default: 120.0)
    pub router_ttl_secs: f64,

166
    /// Maximum tree size before pruning (only used when use_kv_events is false, default: 2^20 = 1048576)
167
168
169
170
    pub router_max_tree_size: usize,

    /// Target size ratio after pruning (only used when use_kv_events is false, default: 0.8)
    pub router_prune_target_ratio: f64,
171
172
173
174
175
}

impl Default for KvRouterConfig {
    fn default() -> Self {
        Self {
176
            overlap_score_weight: 1.0,
177
            router_temperature: 0.0,
178
            use_kv_events: true,
179
            router_replica_sync: false,
180
            router_track_active_blocks: true,
181
            router_track_output_blocks: false,
182
            router_assume_kv_reuse: true,
183
            router_snapshot_threshold: Some(1000000),
184
            router_reset_states: false,
185
            router_ttl_secs: 120.0,
186
            router_max_tree_size: 2usize.pow(20), // 2^20 = 1048576, matches PruneConfig::default()
187
            router_prune_target_ratio: 0.8,
188
189
190
191
192
193
194
        }
    }
}

impl KvRouterConfig {
    /// Create a new KvRouterConfig with optional weight values.
    /// If a weight is None, the default value will be used.
195
    #[allow(clippy::too_many_arguments)]
196
197
    pub fn new(
        overlap_score_weight: Option<f64>,
198
        temperature: Option<f64>,
199
        use_kv_events: Option<bool>,
200
        replica_sync: Option<bool>,
201
        track_active_blocks: Option<bool>,
202
        track_output_blocks: Option<bool>,
203
        assume_kv_reuse: Option<bool>,
204
205
        router_snapshot_threshold: Option<Option<u32>>,
        router_reset_states: Option<bool>,
206
207
208
        router_ttl_secs: Option<f64>,
        router_max_tree_size: Option<usize>,
        router_prune_target_ratio: Option<f64>,
209
210
211
212
    ) -> Self {
        let default = Self::default();
        Self {
            overlap_score_weight: overlap_score_weight.unwrap_or(default.overlap_score_weight),
213
            router_temperature: temperature.unwrap_or(default.router_temperature),
214
            use_kv_events: use_kv_events.unwrap_or(default.use_kv_events),
215
            router_replica_sync: replica_sync.unwrap_or(default.router_replica_sync),
216
217
            router_track_active_blocks: track_active_blocks
                .unwrap_or(default.router_track_active_blocks),
218
219
            router_track_output_blocks: track_output_blocks
                .unwrap_or(default.router_track_output_blocks),
220
            router_assume_kv_reuse: assume_kv_reuse.unwrap_or(default.router_assume_kv_reuse),
221
222
223
            router_snapshot_threshold: router_snapshot_threshold
                .unwrap_or(default.router_snapshot_threshold),
            router_reset_states: router_reset_states.unwrap_or(default.router_reset_states),
224
225
226
227
            router_ttl_secs: router_ttl_secs.unwrap_or(default.router_ttl_secs),
            router_max_tree_size: router_max_tree_size.unwrap_or(default.router_max_tree_size),
            router_prune_target_ratio: router_prune_target_ratio
                .unwrap_or(default.router_prune_target_ratio),
228
229
        }
    }
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260

    /// Compute sequence hashes for active block tracking based on configuration.
    ///
    /// Returns:
    /// - `None` if `router_track_active_blocks` is false
    /// - Random hashes if `router_track_active_blocks` is true but `router_assume_kv_reuse` is false
    /// - Actual sequence hashes if both are true
    pub fn compute_seq_hashes_for_tracking(
        &self,
        tokens: &[u32],
        block_size: u32,
    ) -> Option<Vec<u64>> {
        if !self.router_track_active_blocks {
            return None;
        }

        let num_blocks = tokens.len() / block_size as usize;
        if num_blocks == 0 {
            return Some(Vec::new());
        }

        if self.router_assume_kv_reuse {
            // Compute actual block hashes and sequence hashes
            let block_hashes = compute_block_hash_for_seq(tokens, block_size, None);
            Some(compute_seq_hash_for_block(&block_hashes))
        } else {
            // Generate random hashes (no KV reuse assumed)
            let mut rng = rand::rng();
            Some((0..num_blocks).map(|_| rng.random::<u64>()).collect())
        }
    }
261
262
}

263
pub enum Indexer {
264
265
    /// Updates itself based on KV events emitted by backend workers or routing decisions.
    /// Supports TTL-based expiration and size-based pruning.
266
    /// Has the ability to persist and snapshot states.
267
    KvIndexer(KvIndexer),
268
269
270
271

    /// Used when we do not wish to use the indexer at all (e.g., when overlap_score_weight is 0).
    /// Note: This will cause KV events to accumulate in JetStream as we do not regularly purge them.
    None,
272
273
274
275
276
277
278
279
280
}

impl Indexer {
    async fn find_matches(
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => indexer.find_matches(sequence).await,
281
282
283
            Indexer::None => Ok(OverlapScores {
                scores: HashMap::new(),
                frequencies: Vec::new(),
284
                tree_sizes: HashMap::new(),
285
            }),
286
287
        }
    }
288
289
290
291

    async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => indexer.dump_events().await,
292
293
294
295
296
            Indexer::None => {
                panic!(
                    "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)"
                );
            }
297
298
        }
    }
299

300
    async fn process_routing_decision_for_request(
301
        &self,
302
        tokens_with_hashes: &mut TokensWithHashes,
303
304
305
306
307
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => {
                indexer
308
                    .process_routing_decision_for_request(tokens_with_hashes, worker)
309
310
311
312
313
                    .await
            }
            Indexer::None => Ok(()),
        }
    }
314
315
}

316
317
/// A KvRouter only decides which worker you should use. It doesn't send you there.
/// TODO: Rename this to indicate it only selects a worker, it does not route.
318
pub struct KvRouter {
319
320
321
    indexer: Indexer,

    // How about a Box<dyn KvIndexerInterface>
322
    scheduler: KvScheduler,
323

324
    block_size: u32,
325
326

    kv_router_config: KvRouterConfig,
Yan Ru Pei's avatar
Yan Ru Pei committed
327
328

    cancellation_token: tokio_util::sync::CancellationToken,
329
330

    client: Client,
331
332

    worker_query_client: Option<WorkerQueryClient>,
333
334
335
336
}

impl KvRouter {
    pub async fn new(
337
338
        endpoint: Endpoint,
        client: Client,
339
        workers_with_configs: Arc<RuntimeConfigs>,
340
        block_size: u32,
341
        selector: Option<Box<dyn WorkerSelector + Send + Sync>>,
342
        kv_router_config: Option<KvRouterConfig>,
343
        router_id: u64,
344
    ) -> Result<Self> {
345
        let kv_router_config = kv_router_config.unwrap_or_default();
346
        let component = endpoint.component();
347
        let cancellation_token = component.drt().primary_token();
348

349
350
351
        let indexer = if kv_router_config.overlap_score_weight == 0.0 {
            // When overlap_score_weight is zero, we don't need to track prefixes
            Indexer::None
352
        } else {
353
            let kv_indexer_metrics = indexer::KvIndexerMetrics::from_component(component);
354
355
356
357
358
359
360
361
362
363
364
365
366

            // If use_kv_events is false, enable TTL and pruning for approximate behavior
            let prune_config = if !kv_router_config.use_kv_events {
                Some(PruneConfig {
                    ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs),
                    max_tree_size: kv_router_config.router_max_tree_size,
                    prune_target_ratio: kv_router_config.router_prune_target_ratio,
                })
            } else {
                None
            };

            Indexer::KvIndexer(KvIndexer::new_with_frequency(
367
                cancellation_token.clone(),
368
                None, // expiration_duration for frequency tracking
369
370
                block_size,
                kv_indexer_metrics,
371
                prune_config,
372
373
            ))
        };
374

375
376
377
        // Wait for at least one worker with a known runtime config before starting scheduler
        workers_with_configs.subscribe().wait_for_some().await;

378
        let scheduler = KvScheduler::start(
379
            component.clone(),
380
            block_size,
381
            workers_with_configs.clone(),
382
            selector,
383
            kv_router_config.router_replica_sync,
384
            router_id,
385
386
        )
        .await?;
387

388
389
        // Initialize worker query client using namespace abstraction
        // (created before background task so we can use it for startup recovery)
390
391
392
393
394
        // Uses a subscriber from workers_with_configs
        let worker_query_client = worker_query::WorkerQueryClient::new(
            component.clone(),
            workers_with_configs.subscribe(),
        );
395
396
        tracing::info!("Worker query client initialized");

397
398
399
400
        // Start KV event subscriber background process (only when use_kv_events is enabled)
        if kv_router_config.use_kv_events
            && let Indexer::KvIndexer(ref kv_indexer) = indexer
        {
401
402
403
404
405
406
            let all_local_indexer = workers_with_configs
                .configs
                .iter()
                .filter_map(|r| r.value().as_ref().map(|c| c.enable_local_indexer))
                .all(|b| b);

407
408
409
410
            tracing::info!(
                "Found {} worker(s), starting KV event subscriber",
                workers_with_configs.num_workers()
            );
411

412
413
            let transport_kind = EventTransportKind::from_env_or_default();

414
            // Start subscriber - setup runs synchronously, then spawns background loop internally
415
            if all_local_indexer {
416
417
418
419
420
421
422
423
424
425
                if transport_kind == EventTransportKind::Zmq {
                    if kv_router_config.router_snapshot_threshold.is_some()
                        || kv_router_config.router_reset_states
                    {
                        tracing::warn!(
                            "ZMQ event plane does not support KV snapshots or state reset; ignoring snapshot/reset settings"
                        );
                    }
                } else {
                    tracing::info!(
426
427
                        "All {} workers have local_indexer enabled, using NATS Core subscription",
                        workers_with_configs.num_workers()
428
429
430
431
                    );
                }

                start_kv_router_background_event_plane(
432
433
434
435
436
437
                    component.clone(),
                    kv_indexer.event_sender(),
                    kv_indexer.remove_worker_sender(),
                    cancellation_token.clone(),
                    worker_query::WorkerQueryClient::new(
                        component.clone(),
438
                        workers_with_configs.subscribe(),
439
                    ),
440
                    transport_kind,
441
442
                )
                .await?;
443
            } else {
444
445
446
447
448
                if transport_kind == EventTransportKind::Zmq {
                    tracing::warn!(
                        "Not all workers have local_indexer enabled; falling back to JetStream for durability"
                    );
                }
449
450
451
                tracing::info!(
                    "Not all workers have local_indexer enabled, using JetStream subscription"
                );
452

453
454
                // Convert router_id to string for NATS consumer naming
                let consumer_id = router_id.to_string();
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
                start_kv_router_background(
                    component.clone(),
                    consumer_id,
                    kv_indexer.event_sender(),
                    kv_indexer.remove_worker_sender(),
                    kv_router_config
                        .router_snapshot_threshold
                        .map(|_| kv_indexer.get_workers_sender()),
                    kv_router_config
                        .router_snapshot_threshold
                        .map(|_| kv_indexer.snapshot_event_sender()),
                    cancellation_token.clone(),
                    kv_router_config.router_snapshot_threshold,
                    kv_router_config.router_reset_states,
                )
                .await?;
471
            }
472
        }
473

474
        tracing::info!("KV Routing initialized");
475
        Ok(Self {
476
            indexer,
477
            scheduler,
478
            block_size,
479
            kv_router_config,
Yan Ru Pei's avatar
Yan Ru Pei committed
480
            cancellation_token,
481
            client,
482
            worker_query_client: Some(worker_query_client),
483
        })
484
485
    }

486
487
488
489
490
    /// Get a reference to the client used by this KvRouter
    pub fn client(&self) -> &Client {
        &self.client
    }

491
    /// Give these tokens, find the worker with the best match in it's KV cache.
Yan Ru Pei's avatar
Yan Ru Pei committed
492
    /// Returns the best worker (with dp_rank) and overlap amount in number of blocks.
Yan Ru Pei's avatar
Yan Ru Pei committed
493
    /// Now also takes optional context_id for request tracking
494
    #[allow(clippy::too_many_arguments)]
Yan Ru Pei's avatar
Yan Ru Pei committed
495
    pub async fn find_best_match(
496
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
497
        context_id: Option<&str>,
498
        tokens: &[u32],
499
        router_config_override: Option<&RouterConfigOverride>,
500
        update_states: bool,
501
        lora_name: Option<String>,
Yan Ru Pei's avatar
Yan Ru Pei committed
502
    ) -> anyhow::Result<(WorkerWithDpRank, u32)> {
Yan Ru Pei's avatar
Yan Ru Pei committed
503
504
505
506
507
        // Validate that context_id is provided when update_states is true
        if update_states && context_id.is_none() {
            panic!("context_id must be provided if update_states is true");
        }

508
        let isl_tokens = tokens.len();
509

510
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None);
511
        let overlap_scores = self.indexer.find_matches(block_hashes).await?;
512

513
514
515
        // Compute seq_hashes only if scheduler needs it for active blocks tracking
        let maybe_seq_hashes = self
            .kv_router_config
516
            .compute_seq_hashes_for_tracking(tokens, self.block_size);
517

Yan Ru Pei's avatar
Yan Ru Pei committed
518
        let best_worker = self
519
            .scheduler
520
            .schedule(
Yan Ru Pei's avatar
Yan Ru Pei committed
521
                context_id.map(|s| s.to_string()),
522
                isl_tokens,
523
                maybe_seq_hashes,
524
                overlap_scores.clone(),
525
                router_config_override,
526
                update_states,
527
                lora_name,
528
            )
529
            .await?;
530

531
532
        // Note: Routing decision recording (for approximate mode) is now handled
        // by KvPushRouter::generate after select_worker returns.
533

534
535
        let overlap_amount = overlap_scores
            .scores
Yan Ru Pei's avatar
Yan Ru Pei committed
536
            .get(&best_worker)
537
538
            .copied()
            .unwrap_or(0);
Yan Ru Pei's avatar
Yan Ru Pei committed
539
        Ok((best_worker, overlap_amount))
540
541
    }

542
    #[allow(clippy::too_many_arguments)]
543
544
545
546
547
    pub async fn add_request(
        &self,
        request_id: String,
        tokens: &[u32],
        overlap_blocks: u32,
548
        expected_output_tokens: Option<u32>,
Yan Ru Pei's avatar
Yan Ru Pei committed
549
        worker: WorkerWithDpRank,
550
        lora_name: Option<String>,
551
552
    ) {
        let isl_tokens = tokens.len();
553

554
555
556
        let maybe_seq_hashes = self
            .kv_router_config
            .compute_seq_hashes_for_tracking(tokens, self.block_size);
557

558
559
        if let Err(e) = self
            .scheduler
560
            .add_request(
561
                request_id.clone(),
562
                maybe_seq_hashes,
563
564
                isl_tokens,
                overlap_blocks,
565
                expected_output_tokens,
Yan Ru Pei's avatar
Yan Ru Pei committed
566
                worker,
567
                lora_name,
568
            )
569
570
571
572
            .await
        {
            tracing::warn!("Failed to add request {request_id}: {e}");
        }
573
574
    }

575
    pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> {
576
        self.scheduler.mark_prefill_completed(request_id).await
577
578
    }

579
    pub async fn free(&self, request_id: &str) -> Result<(), SequenceError> {
580
        self.scheduler.free(request_id).await
581
    }
582

583
584
585
586
587
588
589
590
591
592
    pub async fn add_output_block(
        &self,
        request_id: &str,
        decay_fraction: Option<f64>,
    ) -> Result<(), SequenceError> {
        self.scheduler
            .add_output_block(request_id, decay_fraction)
            .await
    }

593
    pub fn block_size(&self) -> u32 {
594
595
        self.block_size
    }
596

597
598
599
600
601
602
603
604
605
606
607
608
    /// Compute the overlap blocks for a given token sequence and worker.
    /// This queries the indexer to find how many blocks are already cached.
    pub async fn get_overlap_blocks(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
    ) -> Result<u32, KvRouterError> {
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None);
        let overlap_scores = self.indexer.find_matches(block_hashes).await?;
        Ok(overlap_scores.scores.get(&worker).copied().unwrap_or(0))
    }

609
610
611
    /// Get potential prefill and decode loads for all workers
    pub async fn get_potential_loads(&self, tokens: &[u32]) -> Result<Vec<PotentialLoad>> {
        let isl_tokens = tokens.len();
612
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None);
613
        let overlap_scores = self.indexer.find_matches(block_hashes.clone()).await?;
614

615
616
        let maybe_seq_hashes = self
            .kv_router_config
617
            .compute_seq_hashes_for_tracking(tokens, self.block_size);
618

619
620
        Ok(self
            .scheduler
621
            .get_potential_loads(maybe_seq_hashes, isl_tokens, overlap_scores)
622
623
624
            .await)
    }

625
626
627
628
    /// Dump all events from the indexer
    pub async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        self.indexer.dump_events().await
    }
629
630
631
632
633
634

    /// Query a specific worker's local KV indexer for its events
    /// (See docstring for `WorkerQueryClient.query_worker()`)
    pub async fn query_worker_local_kv(
        &self,
        worker_id: WorkerId,
635
        dp_rank: DpRank,
636
637
638
639
640
641
642
643
644
        start_event_id: Option<u64>,
        end_event_id: Option<u64>,
    ) -> Result<WorkerKvQueryResponse> {
        let query_client = self
            .worker_query_client
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Worker query client not available (NATS required)"))?;

        query_client
645
            .query_worker(worker_id, dp_rank, start_event_id, end_event_id)
646
647
648
            .await
    }

649
    /// Recover missed KV events from a specific worker's dp_rank.
650
651
652
653
654
655
656
    ///
    /// Queries the worker's local KV indexer for events starting from
    /// `start_event_id` and applies them to the router's indexer.
    ///
    /// # Arguments
    ///
    /// * `worker_id` - The worker to recover from
657
    /// * `dp_rank` - The data parallel rank to recover from
658
659
660
661
662
    /// * `start_event_id` - First event ID to fetch (inclusive), or None to start from beginning
    /// * `end_event_id` - Last event ID to fetch (inclusive), or None for all
    pub async fn recover_from_worker(
        &self,
        worker_id: WorkerId,
663
        dp_rank: DpRank,
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
        start_event_id: Option<u64>,
        end_event_id: Option<u64>,
    ) -> Result<usize> {
        let query_client = self
            .worker_query_client
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Worker query client not available"))?;

        let event_tx = match &self.indexer {
            Indexer::KvIndexer(kv_indexer) => kv_indexer.event_sender(),
            Indexer::None => {
                anyhow::bail!("Cannot recover: indexer is disabled (--overlap_score_weight is 0)")
            }
        };

679
680
681
        query_client
            .recover_from_worker(worker_id, dp_rank, start_event_id, end_event_id, &event_tx)
            .await
682
    }
683
684
}

Michael Feil's avatar
Michael Feil committed
685
686
// NOTE: KVRouter works like a PushRouter,
// but without the reverse proxy functionality, but based on contract of 3 request types
687
688
689
690
691
692
693
#[async_trait]
impl AsyncEngine<SingleIn<RouterRequest>, ManyOut<Annotated<RouterResponse>>, Error> for KvRouter {
    async fn generate(
        &self,
        request: SingleIn<RouterRequest>,
    ) -> Result<ManyOut<Annotated<RouterResponse>>> {
        let (request, ctx) = request.into_parts();
Michael Feil's avatar
Michael Feil committed
694
695
696
        let context_id = ctx.context().id().to_string();
        // Handle different request types
        let response = match request {
697
            RouterRequest::New { tokens } => {
Yan Ru Pei's avatar
Yan Ru Pei committed
698
                let (best_worker, overlap_blocks) = self
699
                    .find_best_match(Some(&context_id), &tokens, None, true, None)
Michael Feil's avatar
Michael Feil committed
700
701
702
                    .await?;

                RouterResponse::New {
Yan Ru Pei's avatar
Yan Ru Pei committed
703
704
                    worker_id: best_worker.worker_id,
                    dp_rank: best_worker.dp_rank,
Michael Feil's avatar
Michael Feil committed
705
706
707
                    overlap_blocks,
                }
            }
708
709
710
711
712
713
            RouterRequest::MarkPrefill => RouterResponse::PrefillMarked {
                success: self.mark_prefill_completed(&context_id).await.is_ok(),
            },
            RouterRequest::MarkFree => RouterResponse::FreeMarked {
                success: self.free(&context_id).await.is_ok(),
            },
Michael Feil's avatar
Michael Feil committed
714
        };
715
716
717
718
719
720

        let response = Annotated::from_data(response);
        let stream = stream::iter(vec![response]);
        Ok(ResponseStream::new(Box::pin(stream), ctx.context()))
    }
}
721
722

pub struct KvPushRouter {
723
    inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
724
    pub chooser: Arc<KvRouter>,
725
726
}

727
728
729
730
731
732
733
/// Result of worker selection containing instance ID, dp_rank, and overlap amount.
struct WorkerSelection {
    instance_id: u64,
    dp_rank: u32,
    overlap_amount: u32,
}

734
735
impl KvPushRouter {
    pub fn new(
736
        inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
737
738
739
740
        chooser: Arc<KvRouter>,
    ) -> Self {
        KvPushRouter { inner, chooser }
    }
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755

    /// Select a worker for the request, either using a preselected worker or finding the best match.
    ///
    /// When `is_query_only` is false and `handle_local_updates` is true, this also registers
    /// the request with the scheduler via `add_request`.
    async fn select_worker(
        &self,
        context_id: &str,
        request: &PreprocessedRequest,
        phase: RequestPhase,
        is_query_only: bool,
        handle_local_updates: bool,
    ) -> Result<WorkerSelection, Error> {
        let routing = request.routing.as_ref();

756
757
758
        // Extract LORA name from routing hints
        let lora_name = routing.and_then(|r| r.lora_name.clone());

759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
        // Get pre-selected worker based on phase, with backend_instance_id as fallback
        let Some(id) = (match phase {
            RequestPhase::Prefill => {
                routing.and_then(|r| r.prefill_worker_id.or(r.backend_instance_id))
            }
            RequestPhase::Decode => {
                routing.and_then(|r| r.decode_worker_id.or(r.backend_instance_id))
            }
            RequestPhase::Aggregated => routing.and_then(|r| r.backend_instance_id),
        }) else {
            // No preselected worker - find the best match
            // Don't update states if this is a query-only request
            let (best_worker, overlap_amount) = self
                .chooser
                .find_best_match(
                    Some(context_id),
                    &request.token_ids,
                    request.router_config_override.as_ref(),
                    !is_query_only,
778
                    lora_name,
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
                )
                .await?;

            return Ok(WorkerSelection {
                instance_id: best_worker.worker_id,
                dp_rank: best_worker.dp_rank,
                overlap_amount,
            });
        };

        // Route to pre-selected or explicitly specified worker
        let dp_rank = routing.and_then(|r| r.dp_rank).unwrap_or(0);
        tracing::debug!(
            worker_id = id,
            dp_rank = dp_rank,
            ?phase,
            "Routing to specified worker"
        );

        // Compute actual overlap blocks by querying the indexer
        let worker = WorkerWithDpRank::new(id, dp_rank);
        let overlap_blocks = self
            .chooser
            .get_overlap_blocks(&request.token_ids, worker)
            .await?;

805
806
807
808
809
810
        // Extract expected_output_tokens from routing hints
        let expected_output_tokens = request
            .routing
            .as_ref()
            .and_then(|r| r.expected_output_tokens);

811
812
813
814
815
816
817
        // Perform add_request if this router handles local updates
        if !is_query_only && handle_local_updates {
            self.chooser
                .add_request(
                    context_id.to_string(),
                    &request.token_ids,
                    overlap_blocks,
818
                    expected_output_tokens,
819
                    worker,
820
                    lora_name,
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
                )
                .await;
        } else {
            tracing::debug!(
                request_id = %context_id,
                worker_id = id,
                dp_rank = dp_rank,
                "Skipping add_request - query or handled externally"
            );
        }

        Ok(WorkerSelection {
            instance_id: id,
            dp_rank,
            overlap_amount: overlap_blocks,
        })
    }
838
839
840
}

#[async_trait]
841
impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<Annotated<LLMEngineOutput>>, Error>
842
843
    for KvPushRouter
{
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
    /// Generate method that handles KV-aware routing with three distinct behaviors:
    ///
    /// 1. **If `query_instance_id` annotation is set**:
    ///    - Returns the best matching worker ID without routing the request
    ///    - Does NOT update any router local states
    ///    - Response includes worker_instance_id and token_data annotations
    ///
    /// 2. **If `backend_instance_id` is set in the request**:
    ///    - Routes directly to the specified backend instance
    ///    - DOES update router states to track this request (unless query_instance_id is also set)
    ///    - Bypasses the normal KV matching logic
    ///
    /// 3. **If neither are set (default behavior)**:
    ///    - Finds the best worker based on KV cache overlap
    ///    - Updates router states to track the request
    ///    - Routes to the selected worker
    ///
    /// The router state updates include tracking active sequences and managing
    /// prefill/completion lifecycle for proper KV cache management.
863
864
    async fn generate(
        &self,
865
        request: SingleIn<PreprocessedRequest>,
866
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
867
868
869
        // Extract context ID for request tracking
        let context_id = request.context().id().to_string();

870
871
872
        // Simple query-only detection: presence of query_instance_id annotation means query-only mode
        let is_query_only = request.get_annotation_value("query_instance_id").is_some();

873
        // Determine if this router should handle local state updates (add_request, free, etc.)
874
875
876
877
878
879
        // Default is true (router handles bookkeeping). Set to false for GAIE Stage 2 where
        // an external orchestrator (e.g., EPP sidecar) handles bookkeeping via C FFI.
        let handle_local_updates = request
            .routing
            .as_ref()
            .and_then(|r| r.enable_local_updates)
880
881
            .unwrap_or(true);

882
883
884
        // Get phase from tracker (defaults to Aggregated if no tracker or phase not set)
        let phase = request
            .tracker
885
            .as_ref()
886
887
888
889
            .map(|t| t.phase())
            .unwrap_or(RequestPhase::Aggregated);

        let block_size = self.chooser.block_size() as usize;
890
891
892
893
894
895
896
897
898
899
900
901
902
903
        let selection = self
            .select_worker(
                &context_id,
                &request,
                phase,
                is_query_only,
                handle_local_updates,
            )
            .await?;
        let WorkerSelection {
            instance_id,
            dp_rank,
            overlap_amount,
        } = selection;
904

905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
        // In approximate mode (use_kv_events=false), record the routing decision
        // so the indexer can track cache state based on routing decisions.
        // This covers both pre-selected workers and find_best_match selections.
        if !is_query_only && !self.chooser.kv_router_config.use_kv_events {
            let worker = WorkerWithDpRank::new(instance_id, dp_rank);
            let mut tokens_with_hashes =
                TokensWithHashes::new(request.token_ids.clone(), self.chooser.block_size);
            if let Err(e) = self
                .chooser
                .indexer
                .process_routing_decision_for_request(&mut tokens_with_hashes, worker)
                .await
            {
                tracing::warn!(
                    request_id = %context_id,
                    worker_id = instance_id,
                    dp_rank = dp_rank,
                    error = %e,
                    "Failed to record routing decision in approximate mode"
                );
            }
        }

928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
        // Record metrics in tracker: KV hit rate and worker ID based on phase
        if let Some(ref tracker) = request.tracker {
            let isl_blocks = request.token_ids.len().div_ceil(block_size);
            tracker.record_kv_hit(overlap_amount, isl_blocks);
            tracker.record_worker(instance_id);
        }

        // Handle query-only requests: early return with worker info
        if is_query_only {
            let stream_context = request.context().clone();
            // Tracker is always created for query-only requests (delta generator enables tracking
            // when query_instance_id annotation is present)
            let worker_id_info = request.tracker.as_ref().and_then(|t| t.get_worker_info());

            tracing::trace!(
                ?phase,
                worker_id = instance_id,
                ?worker_id_info,
                "Returning worker selection (query-only mode)"
            );

949
950
951
952
953
954
955
956
957
            let output = LLMEngineOutput {
                disaggregated_params: Some(json!({
                    "worker_id": worker_id_info,
                    "token_ids": request.token_ids
                })),
                ..Default::default()
            };
            let response = Annotated::from_data(output);
            let stream = stream::iter(vec![response]);
958
959
            return Ok(ResponseStream::new(Box::pin(stream), stream_context));
        }
960
961

        // Route to worker
962
963
964
965
966
967
968
969
        let isl_tokens = request.token_ids.len();
        let expected_output_tokens = request
            .routing
            .as_ref()
            .and_then(|r| r.expected_output_tokens);
        let track_output_blocks =
            self.chooser.kv_router_config.router_track_output_blocks && handle_local_updates;

970
        let (mut backend_input, context) = request.into_parts();
971
        backend_input.routing_mut().dp_rank = Some(dp_rank);
972
973
        let updated_request = context.map(|_| backend_input);

974
        let chooser = self.chooser.clone();
975
976
977
978
        let mut response_stream = self.inner.direct(updated_request, instance_id).await?;
        let stream_context = response_stream.context();
        let context_for_monitoring = stream_context.clone();

979
980
981
        // Wrap stream with lifecycle management (mark_prefill_completed, free)
        // Only perform these operations if handle_local_updates is true.
        // When false, an external caller (e.g., GAIE sidecar) handles bookkeeping via C FFI.
982
983
984
        let wrapped_stream = Box::pin(async_stream::stream! {
            let mut prefill_marked = false;

985
986
987
988
            // Output block tracking state
            let mut cumulative_osl: usize = 0;
            let mut current_total_blocks = isl_tokens.div_ceil(block_size);

989
990
991
992
993
994
995
            loop {
                tokio::select! {
                    biased;

                    _ = context_for_monitoring.stopped() => {
                        tracing::debug!("Request {context_id} cancelled, ending stream");
                        break;
996
                    }
Yan Ru Pei's avatar
Yan Ru Pei committed
997

998
                    item = response_stream.next() => {
999
                        let Some(item) = item else {
1000
1001
                            break;
                        };
1002

1003
                        if handle_local_updates && !prefill_marked {
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
                            // Only mark prefill completed when we receive actual tokens,
                            // not empty bootstrap info (token_ids: []) from disaggregated prefill
                            let has_tokens = item.data.as_ref()
                                .map(|d| !d.token_ids.is_empty())
                                .unwrap_or(false);
                            if has_tokens {
                                if let Err(e) = chooser.mark_prefill_completed(&context_id).await {
                                    tracing::warn!("Failed to mark prefill completed for request {context_id}: {e}");
                                }
                                prefill_marked = true;
1014
                            }
1015
                        }
1016

1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
                        // Track output blocks if enabled
                        if track_output_blocks {
                            let new_tokens = item.data.as_ref()
                                .map(|d| d.token_ids.len())
                                .unwrap_or(0);
                            cumulative_osl += new_tokens;

                            let new_total_blocks = (isl_tokens + cumulative_osl).div_ceil(block_size);
                            if new_total_blocks > current_total_blocks {
                                // New block boundary crossed - add output block with decay
                                // Clamp eot to min 1 to avoid division by zero, and result to min 0.0
                                let decay_fraction = expected_output_tokens.map(|eot| {
                                    (1.0 - (cumulative_osl as f64 / eot.max(1) as f64)).max(0.0)
                                });
                                if let Err(e) = chooser.add_output_block(&context_id, decay_fraction).await {
                                    tracing::warn!(
                                        "Failed to add output block for request {context_id}: {e}"
                                    );
                                }
                                current_total_blocks = new_total_blocks;
                            }
                        }

1040
                        yield item;
1041
                    }
1042
1043
                }
            }
1044

1045
1046
1047
1048
1049
            // Only call free() if we handle local updates.
            // When handle_local_updates=false, external caller handles cleanup via C FFI.
            if handle_local_updates
                && let Err(e) = chooser.free(&context_id).await
            {
1050
                tracing::warn!("Failed to free request {context_id}: {e}");
1051
            }
1052
1053
        });
        Ok(ResponseStream::new(wrapped_stream, stream_context))
1054
1055
    }
}
Yan Ru Pei's avatar
Yan Ru Pei committed
1056
1057
1058
1059
1060
1061
1062

impl Drop for KvRouter {
    fn drop(&mut self) {
        tracing::info!("Dropping KvRouter - cancelling background tasks");
        self.cancellation_token.cancel();
    }
}