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

use std::collections::HashMap;

use rand::Rng;
7
use rustc_hash::FxHashMap;
8
9

use super::config::KvRouterConfig;
10
use super::types::{KvSchedulerError, SchedulingRequest, pinned_worker_config};
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
use crate::protocols::{WorkerConfigLike, WorkerId, WorkerSelectionResult, WorkerWithDpRank};

/// A trait that users can implement to define custom selection logic.
///
/// Generic over `C` so that the scheduling layer does not depend on a concrete config type.
pub trait WorkerSelector<C: WorkerConfigLike> {
    fn select_worker(
        &self,
        workers: &HashMap<WorkerId, C>,
        request: &SchedulingRequest,
        block_size: u32,
    ) -> Result<WorkerSelectionResult, KvSchedulerError>;
}

/// Helper function for softmax sampling.
26
/// Returns the selected worker and its logit.
27
fn softmax_sample(
28
    logits: &FxHashMap<WorkerWithDpRank, f64>,
29
    temperature: f64,
30
31
32
33
34
35
) -> (WorkerWithDpRank, f64) {
    let mut rng = rand::rng();
    softmax_sample_with_sample(logits, temperature, rng.random())
}

fn softmax_sample_with_sample(
36
    logits: &FxHashMap<WorkerWithDpRank, f64>,
37
38
39
    temperature: f64,
    sample: f64,
) -> (WorkerWithDpRank, f64) {
40
    assert!(!logits.is_empty(), "Empty logits for softmax sampling");
41
42

    if temperature == 0.0 {
43
44
45
46
47
        let (worker, logit) = logits
            .iter()
            .min_by(|a, b| a.1.total_cmp(b.1))
            .expect("logits non-empty");
        return (*worker, *logit);
48
49
    }

50
    let entries: Vec<(WorkerWithDpRank, f64)> = logits.iter().map(|(w, l)| (*w, *l)).collect();
51

52
53
54
55
56
    let (min_val, max_val) = entries
        .iter()
        .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), (_, v)| {
            (lo.min(*v), hi.max(*v))
        });
57

58
    let mut probs = if min_val == max_val {
59
        vec![1.0 / entries.len() as f64; entries.len()]
60
    } else {
61
62
63
64
65
66
67
68
        // Negate logits and rescale to [−1/temperature, 0] for numerical stability
        // before softmax. Subtracting the max (which maps to min_val) keeps exp() inputs ≤ 0.
        let scale = -1.0 / ((max_val - min_val) * temperature);
        let max_scaled = min_val * scale;
        entries
            .iter()
            .map(|(_, v)| (v * scale - max_scaled).exp())
            .collect::<Vec<f64>>()
69
70
    };

71
72
73
    let sum: f64 = probs.iter().sum();
    probs.iter_mut().for_each(|p| *p /= sum);

74
    let mut cumsum = 0.0;
75
    for (i, &prob) in probs.iter().enumerate() {
76
77
        cumsum += prob;
        if sample <= cumsum {
78
            return entries[i];
79
80
81
        }
    }

82
    *entries.last().unwrap()
83
84
85
}

/// Default implementation matching the Python _cost_function.
86
#[derive(Debug, Clone)]
87
88
pub struct DefaultWorkerSelector {
    pub kv_router_config: KvRouterConfig,
89
    pub worker_type: &'static str,
90
91
92
}

impl DefaultWorkerSelector {
93
    pub fn new(kv_router_config: Option<KvRouterConfig>, worker_type: &'static str) -> Self {
94
95
        Self {
            kv_router_config: kv_router_config.unwrap_or_default(),
96
            worker_type,
97
98
        }
    }
99

100
    fn worker_logit(
101
102
103
104
105
        &self,
        request: &SchedulingRequest,
        worker: WorkerWithDpRank,
        block_size: u32,
        overlap_weight: f64,
106
        shared_cache_multiplier: f64,
107
        formula_name: &'static str,
108
    ) -> f64 {
109
        let isl = request.isl_tokens;
110
111
112
113
114
115
116
117
        let effective_overlap_blocks = request
            .effective_overlap_blocks
            .get(&worker)
            .copied()
            .unwrap_or(0.0);
        // `shared_cache_hits::hits_beyond` expects an integer block count, so
        // round the weighted overlap for this comparison only.
        let device_overlap_blocks = effective_overlap_blocks.round().max(0.0) as u32;
118
119
120
121
122
123
        let default_prefill_token = if request.track_prefill_tokens { isl } else { 0 };
        let prefill_token = request
            .prefill_tokens
            .get(&worker)
            .copied()
            .unwrap_or(default_prefill_token);
124
125
126
127

        // Adjust prefill tokens by shared cache hits beyond this worker's device prefix.
        let (adjusted_prefill_token, shared_beyond) =
            if let Some(ref shared_hits) = request.shared_cache_hits {
128
                let beyond = shared_hits.hits_beyond(device_overlap_blocks);
129
130
131
132
133
134
135
136
137
                let reduction = shared_cache_multiplier * (beyond as f64) * (block_size as f64);
                let adjusted = (prefill_token as f64 - reduction).max(0.0) as usize;
                (adjusted, beyond)
            } else {
                (prefill_token, 0)
            };

        let potential_prefill_block = (adjusted_prefill_token as f64) / (block_size as f64);
        let decode_block_fallback = (prefill_token as f64) / (block_size as f64);
138
139
140
141
        let decode_block = request
            .decode_blocks
            .get(&worker)
            .copied()
142
            .unwrap_or(decode_block_fallback.floor() as usize) as f64;
143
144
        let logit = overlap_weight * potential_prefill_block + decode_block;

145
146
        if shared_beyond > 0 {
            tracing::debug!(
147
                "{formula_name} for worker_id={} dp_rank={:?} with {effective_overlap_blocks:.2} effective device blocks, \
148
149
150
151
152
153
154
155
156
                 {shared_beyond} shared blocks beyond device (multiplier={shared_cache_multiplier:.2}): {logit:.3} \
                 = {overlap_weight:.1} * adjusted_prefill_blocks + decode_blocks \
                 = {overlap_weight:.1} * {potential_prefill_block:.3} + {decode_block:.3} \
                 (prefill_tokens: {prefill_token} -> {adjusted_prefill_token})",
                worker.worker_id,
                worker.dp_rank
            );
        } else {
            tracing::debug!(
157
                "{formula_name} for worker_id={} dp_rank={:?} with {effective_overlap_blocks:.2} effective cached blocks: {logit:.3} \
158
159
160
161
162
163
                 = {overlap_weight:.1} * prefill_blocks + decode_blocks \
                 = {overlap_weight:.1} * {potential_prefill_block:.3} + {decode_block:.3}",
                worker.worker_id,
                worker.dp_rank
            );
        }
164

165
        logit
166
    }
167
168
169
170
171
172
173
174
175
176
}

impl<C: WorkerConfigLike> WorkerSelector<C> for DefaultWorkerSelector {
    fn select_worker(
        &self,
        workers: &HashMap<WorkerId, C>,
        request: &SchedulingRequest,
        block_size: u32,
    ) -> Result<WorkerSelectionResult, KvSchedulerError> {
        assert!(request.isl_tokens > 0);
177
        request.validate_worker_constraints()?;
178
179

        let allowed_ids = request.allowed_worker_ids.as_ref();
180
        let pinned_worker = request.pinned_worker;
181

182
183
184
185
186
        if pinned_worker.is_none()
            && allowed_ids.map_or(workers.is_empty(), |ids| {
                !workers.keys().any(|wid| ids.contains(wid))
            })
        {
187
188
189
190
191
192
            return Err(KvSchedulerError::NoEndpoints);
        }

        let isl = request.isl_tokens;
        let request_blocks = isl.div_ceil(block_size as usize);

193
194
195
196
197
198
199
200
201
202
203
204
        let overlap_weight = request
            .router_config_override
            .as_ref()
            .and_then(|cfg| cfg.overlap_score_weight)
            .unwrap_or(self.kv_router_config.overlap_score_weight);

        let shared_cache_multiplier = request
            .router_config_override
            .as_ref()
            .and_then(|cfg| cfg.shared_cache_multiplier)
            .unwrap_or(self.kv_router_config.shared_cache_multiplier);

205
206
207
        if let Some(worker) = pinned_worker {
            pinned_worker_config(workers, worker)?;

208
            let logit = self.worker_logit(
209
210
211
212
                request,
                worker,
                block_size,
                overlap_weight,
213
                shared_cache_multiplier,
214
215
                "Pinned formula",
            );
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
            let effective_overlap_blocks = request
                .effective_overlap_blocks
                .get(&worker)
                .copied()
                .unwrap_or(0.0);
            let cached_tokens = request
                .effective_cached_tokens
                .get(&worker)
                .copied()
                .unwrap_or(0);

            tracing::info!(
                "Selected pinned worker: worker_type={}, worker_id={} dp_rank={:?}, logit: {:.3}, effective cached blocks: {:.2}",
                self.worker_type,
                worker.worker_id,
                worker.dp_rank,
                logit,
                effective_overlap_blocks,
            );
235
236
237
238

            return Ok(WorkerSelectionResult {
                worker,
                required_blocks: request_blocks as u64,
239
240
                effective_overlap_blocks,
                cached_tokens,
241
242
            });
        }
243

244
245
246
247
248
        let temperature = request
            .router_config_override
            .as_ref()
            .and_then(|cfg| cfg.router_temperature)
            .unwrap_or(self.kv_router_config.router_temperature);
249

250
        let get_score = |worker: WorkerWithDpRank| -> f64 {
251
            self.worker_logit(
252
253
254
255
256
257
258
                request,
                worker,
                block_size,
                overlap_weight,
                shared_cache_multiplier,
                "Formula",
            )
259
        };
260

261
262
263
264
265
266
267
268
269
        let worker_iter = workers
            .iter()
            .filter(move |(wid, _)| allowed_ids.is_none_or(|ids| ids.contains(wid)))
            .flat_map(|(worker_id, config)| {
                let data_parallel_size = config.data_parallel_size();
                let data_parallel_start_rank = config.data_parallel_start_rank();
                (data_parallel_start_rank..(data_parallel_start_rank + data_parallel_size))
                    .map(move |dp_rank| WorkerWithDpRank::new(*worker_id, dp_rank))
            });
270

271
272
273
274
275
276
277
278
279
280
281
282
283
        let (best_worker, best_logit) = if temperature == 0.0 {
            let mut min_workers = Vec::new();
            let mut min_score = f64::INFINITY;
            for worker in worker_iter {
                let score = get_score(worker);
                if score < min_score {
                    min_workers.clear();
                    min_workers.push(worker);
                    min_score = score;
                } else if score == min_score {
                    min_workers.push(worker);
                }
            }
284

285
            if min_workers.len() > 1 {
286
                tracing::debug!(
287
                    "Multiple workers tied with same logit, using tree size as tie-breaker"
288
                );
289
290
                let tree_sizes: Vec<(usize, &WorkerWithDpRank)> = min_workers
                    .iter()
291
                    .map(|w| (request.tree_sizes.get(w).copied().unwrap_or(0), w))
292
293
294
295
296
297
298
299
300
                    .collect();

                if tree_sizes.iter().all(|(s, _)| *s == tree_sizes[0].0) {
                    let idx = rand::rng().random_range(0..min_workers.len());
                    (min_workers[idx], min_score)
                } else {
                    let (_, worker) = *tree_sizes.iter().min_by_key(|(s, _)| *s).unwrap();
                    (*worker, min_score)
                }
301
            } else {
302
                (min_workers[0], min_score)
303
304
            }
        } else {
305
            let mut worker_logits = FxHashMap::default();
306
307
308
309
            for worker in worker_iter {
                let score = get_score(worker);
                worker_logits.insert(worker, score);
            }
310

311
312
            softmax_sample(&worker_logits, temperature)
        };
313

314
315
316
317
318
319
320
321
322
323
324
325
326
        let best_host_pinned_overlap_blocks = request
            .tier_overlap_blocks
            .host_pinned
            .get(&best_worker)
            .copied()
            .unwrap_or(0);
        let best_disk_overlap_blocks = request
            .tier_overlap_blocks
            .disk
            .get(&best_worker)
            .copied()
            .unwrap_or(0);

327
328
        if self.worker_type == "decode" {
            tracing::info!(
329
                "Selected worker: worker_type={}, worker_id={} dp_rank={:?}, logit: {:.3}, host_pinned blocks: {}, disk blocks: {}",
330
331
332
333
                self.worker_type,
                best_worker.worker_id,
                best_worker.dp_rank,
                best_logit,
334
335
                best_host_pinned_overlap_blocks,
                best_disk_overlap_blocks,
336
            );
337
338
339
340
341
342
343
344
345
346
347
            let effective_overlap_blocks = request
                .effective_overlap_blocks
                .get(&best_worker)
                .copied()
                .unwrap_or(0.0);
            let cached_tokens = request
                .effective_cached_tokens
                .get(&best_worker)
                .copied()
                .unwrap_or(0);

348
349
350
            return Ok(WorkerSelectionResult {
                worker: best_worker,
                required_blocks: request_blocks as u64,
351
352
                effective_overlap_blocks,
                cached_tokens,
353
354
355
            });
        }

356
357
358
359
360
361
362
363
364
365
        let best_overlap = request
            .effective_overlap_blocks
            .get(&best_worker)
            .copied()
            .unwrap_or(0.0);
        let best_cached_tokens = request
            .effective_cached_tokens
            .get(&best_worker)
            .copied()
            .unwrap_or(0);
366
367
368
369
370
371
372

        let total_blocks_info = workers
            .get(&best_worker.worker_id)
            .and_then(|cfg| cfg.total_kv_blocks())
            .map(|blocks| format!(", total blocks: {}", blocks))
            .unwrap_or_default();

373
        let tree_size = request.tree_sizes.get(&best_worker).copied().unwrap_or(0);
374
375

        tracing::info!(
376
            "Selected worker: worker_type={}, worker_id={} dp_rank={:?}, logit: {:.3}, effective cached blocks: {:.2}, host_pinned blocks: {}, disk blocks: {}, tree size: {}{}",
377
            self.worker_type,
378
379
380
381
            best_worker.worker_id,
            best_worker.dp_rank,
            best_logit,
            best_overlap,
382
383
            best_host_pinned_overlap_blocks,
            best_disk_overlap_blocks,
384
385
386
387
388
389
390
            tree_size,
            total_blocks_info
        );

        Ok(WorkerSelectionResult {
            worker: best_worker,
            required_blocks: request_blocks as u64,
391
392
            effective_overlap_blocks: best_overlap,
            cached_tokens: best_cached_tokens,
393
394
395
396
397
398
399
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
400
    use crate::protocols::SharedCacheHits;
401
402
403

    #[test]
    fn test_softmax_sample_single_key() {
404
        let mut logits = FxHashMap::default();
405
        let worker = WorkerWithDpRank::from_worker_id(42);
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
        for (logit, temperature) in [
            (0.5, 0.1),
            (0.5, 1.0),
            (0.5, 10.0),
            (-100.0, 1.0),
            (100.0, 1.0),
            (0.0, 1.0),
            (0.0, 0.0),
        ] {
            logits.clear();
            logits.insert(worker, logit);

            let result = softmax_sample(&logits, temperature);
            assert_eq!(result.0, worker, "Should return the only available worker");
            assert_eq!(result.1, logit, "Should return the selected worker's logit");
421
422
423
424
425
        }
    }

    #[test]
    fn test_softmax_sample_zero_temperature() {
426
        let mut logits = FxHashMap::default();
427
428
429
430
431
432
433
434
435
436
437
        let worker1 = WorkerWithDpRank::from_worker_id(1);
        let worker2 = WorkerWithDpRank::from_worker_id(2);
        let worker3 = WorkerWithDpRank::from_worker_id(3);
        let worker4 = WorkerWithDpRank::from_worker_id(4);
        logits.insert(worker1, 5.0);
        logits.insert(worker2, 3.0);
        logits.insert(worker3, 7.0);
        logits.insert(worker4, 3.5);

        let result = softmax_sample(&logits, 0.0);
        assert_eq!(
438
439
            result.0, worker2,
            "Should return worker with smallest logit when temperature is 0"
440
441
        );
        assert_eq!(
442
443
            result.1, 3.0,
            "Should return the smallest logit when temperature is 0"
444
445
446
447
448
449
450
451
452
453
454
455
        );

        logits.clear();
        let worker5 = WorkerWithDpRank::from_worker_id(5);
        let worker6 = WorkerWithDpRank::from_worker_id(6);
        logits.insert(worker1, 5.0);
        logits.insert(worker2, 3.0);
        logits.insert(worker5, 3.0);
        logits.insert(worker6, 7.0);

        let result = softmax_sample(&logits, 0.0);
        assert!(
456
457
            result.0 == worker2 || result.0 == worker5,
            "Should return one of the workers tied for the smallest logit"
458
        );
459
        assert_eq!(result.1, 3.0, "Should return the tied minimum logit");
460
461
462
463
464
465
466
467
468
469
470

        logits.clear();
        let worker10 = WorkerWithDpRank::from_worker_id(10);
        let worker20 = WorkerWithDpRank::from_worker_id(20);
        let worker30 = WorkerWithDpRank::from_worker_id(30);
        logits.insert(worker10, -1.0);
        logits.insert(worker20, -5.0);
        logits.insert(worker30, 0.0);

        let result = softmax_sample(&logits, 0.0);
        assert_eq!(
471
            result.0, worker20,
472
473
            "Should handle negative logits correctly"
        );
474
475
476
477
478
479
480
481
482
        assert_eq!(result.1, -5.0, "Should return the minimum negative logit");
    }

    #[test]
    fn test_softmax_sample_with_sample_returns_selected_logit() {
        let worker1 = WorkerWithDpRank::from_worker_id(1);
        let worker2 = WorkerWithDpRank::from_worker_id(2);
        let worker3 = WorkerWithDpRank::from_worker_id(3);

483
        let logits = FxHashMap::from_iter([(worker1, 0.0), (worker2, 3.0), (worker3, 9.0)]);
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
        let entries: Vec<_> = logits
            .iter()
            .map(|(worker, logit)| (*worker, *logit))
            .collect();
        let values: Vec<_> = entries.iter().map(|(_, logit)| *logit).collect();

        let min_val = values.iter().fold(f64::INFINITY, |a, &b| a.min(b));
        let max_val = values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
        let temperature = 1.0;
        let range = max_val - min_val;
        let scaled: Vec<f64> = values.iter().map(|&v| -(v / range) / temperature).collect();
        let max_scaled = scaled.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
        let mut probabilities: Vec<f64> = scaled.iter().map(|&v| (v - max_scaled).exp()).collect();
        let sum: f64 = probabilities.iter().sum();
        probabilities.iter_mut().for_each(|p| *p /= sum);

        let target_idx = entries
            .iter()
            .position(|(_, logit)| *logit > min_val)
            .expect("expected at least one non-minimum logit");
        let cumsum_before: f64 = probabilities.iter().take(target_idx).sum();
        let sample = cumsum_before + probabilities[target_idx] / 2.0;

        let result = softmax_sample_with_sample(&logits, temperature, sample);
        assert_eq!(result, entries[target_idx]);
509
    }
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528

    /// Test the scoring formula with shared cache hits.
    ///
    /// Request [A, B, C, D], shared_cache_multiplier=0.5, block_size=1
    /// - Worker 0: device=[A,B] (overlap=2), shared has [A,B,C,D] -> shared_beyond=2
    ///   adjusted_prefill = isl - 0.5*2*1 = 4-1 = 3, logit = 1.0 * 3 + 0 = 3.0
    /// - Worker 1: device=[] (overlap=0), shared has [A,B,C,D] -> shared_beyond=4
    ///   adjusted_prefill = isl - 0.5*4*1 = 4-2 = 2, logit = 1.0 * 2 + 0 = 2.0
    ///
    /// Worker 1 has lower logit (less work), so it wins.
    #[test]
    fn test_shared_cache_hits_scoring() {
        use crate::test_utils::SimpleWorkerConfig;

        let block_size = 1u32;
        let isl = 4usize;
        let worker0 = WorkerWithDpRank::from_worker_id(0);
        let worker1 = WorkerWithDpRank::from_worker_id(1);

529
530
        let mut effective_overlap_blocks = HashMap::new();
        effective_overlap_blocks.insert(worker0, 2.0);
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
        // worker1 has 0 overlap (not in map)

        #[allow(clippy::single_range_in_vec_init)]
        let shared_hits = SharedCacheHits::from_ranges(vec![0..4]);

        let config = KvRouterConfig {
            overlap_score_weight: 1.0,
            shared_cache_multiplier: 0.5,
            router_temperature: 0.0,
            ..Default::default()
        };

        let selector = DefaultWorkerSelector::new(Some(config), "test");
        let mut workers = HashMap::new();
        workers.insert(0, SimpleWorkerConfig::default());
        workers.insert(1, SimpleWorkerConfig::default());

        let (tx, _rx) = tokio::sync::oneshot::channel();
        let request = SchedulingRequest {
            maybe_request_id: Some("test".into()),
            token_seq: None,
            isl_tokens: isl,
553
554
555
556
            tier_overlap_blocks: Default::default(),
            effective_overlap_blocks,
            effective_cached_tokens: HashMap::new(),
            tree_sizes: HashMap::new(),
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
            decode_blocks: FxHashMap::default(),
            prefill_tokens: FxHashMap::default(),
            track_prefill_tokens: true,
            router_config_override: None,
            update_states: false,
            lora_name: None,
            priority_jump: 0.0,
            expected_output_tokens: None,
            pinned_worker: None,
            allowed_worker_ids: None,
            shared_cache_hits: Some(shared_hits),
            resp_tx: Some(tx),
        };

        let result = selector
            .select_worker(&workers, &request, block_size)
            .unwrap();

        // Worker 1 should win: logit 2.0 < 3.0
        assert_eq!(
            result.worker, worker1,
            "Worker 1 should be selected (lower logit due to shared cache)"
        );
    }

    /// Without shared cache hits, the scoring should be unchanged.
    #[test]
    fn test_no_shared_cache_unchanged() {
        use crate::test_utils::SimpleWorkerConfig;

        let block_size = 16u32;
        let isl = 64usize;
        let worker0 = WorkerWithDpRank::from_worker_id(0);

591
592
        let mut effective_overlap_blocks = HashMap::new();
        effective_overlap_blocks.insert(worker0, 2.0);
593
594
595
596
597
598
599
600
601
602
603

        let config = KvRouterConfig::default();
        let selector = DefaultWorkerSelector::new(Some(config), "test");
        let mut workers = HashMap::new();
        workers.insert(0, SimpleWorkerConfig::default());

        let (tx, _rx) = tokio::sync::oneshot::channel();
        let request = SchedulingRequest {
            maybe_request_id: Some("test".into()),
            token_seq: None,
            isl_tokens: isl,
604
605
606
607
            tier_overlap_blocks: Default::default(),
            effective_overlap_blocks,
            effective_cached_tokens: HashMap::new(),
            tree_sizes: HashMap::new(),
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
            decode_blocks: FxHashMap::default(),
            prefill_tokens: FxHashMap::default(),
            track_prefill_tokens: true,
            router_config_override: None,
            update_states: false,
            lora_name: None,
            priority_jump: 0.0,
            expected_output_tokens: None,
            pinned_worker: None,
            allowed_worker_ids: None,
            shared_cache_hits: None,
            resp_tx: Some(tx),
        };

        let result = selector
            .select_worker(&workers, &request, block_size)
            .unwrap();

        assert_eq!(result.worker, worker0);
    }
628
}