scheduler.rs 34.7 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Asynchronous Scheduler for LLM Request Management
//!
//! This module implements an asynchronous scheduler that handles three main functions:
//! 1. Receiving new requests and placing them in the waiting queue
//! 2. Scheduling waiting requests against available KV cache resources
//! 3. Simulating the execution of running requests with realistic timing
//!
//! ## Scheduling Process
//! The scheduler uses a watermark-based approach to determine if there's sufficient
//! KV cache space for new requests. It also enforces a batched tokens budget to prevent
//! oversubscription of computational resources. Only requests that can be allocated
//! these resources are moved from waiting to running state.
//!
//! ## Request Simulation
//! The simulation models two key phases:
//! - Prefill phase: Uses a quadratic cost function: (cached_tokens + new_tokens) * new_tokens
//! - Decode phase: Uses a cost function proportional to active KV blocks (linear)
//!
//! ## Resource Management
//! The scheduler communicates with the KvManager through MoveBlock signals at each
//! stage of request processing. When resources become constrained, it employs an
//! LRU-based preemption strategy where the oldest running request is evicted and
//! placed at the back of the waiting queue to be rescheduled later.
//!
//! ## NOTE
//! The current prefill and decoding time simulations are not scientific at all and are WIP

31
use crate::kv_router::protocols::{ForwardPassMetrics, KvCacheEventData, KvStats, WorkerStats};
32
33
use crate::mocker::evictor::LRUEvictor;
use crate::mocker::kv_manager::KvManager;
34
use crate::mocker::protocols::{DirectRequest, MockEngineArgs, MoveBlockResponse};
35
use crate::mocker::protocols::{MoveBlock, OutputSignal, PrefillCost, block_response_to_kv_event};
36
use crate::mocker::sequence::ActiveSequence;
37
use crate::tokens::BlockHash;
38
use crate::tokens::blocks::UniqueBlock;
39
40
use std::collections::HashMap;
use std::collections::VecDeque;
41
use tokio::sync::mpsc;
42
use tokio::time::Duration;
43
44
45
46
47
48
49
50
51
52
53
54
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

/// Enum representing either a direct request or an active sequence
pub enum Request {
    Direct(DirectRequest),
    Active(ActiveSequence),
}

#[derive(Default)]
struct SchedulerState {
    waiting: VecDeque<Uuid>,
55
56
    prefill: VecDeque<Uuid>,
    decode: LRUEvictor<Uuid>,
57
    requests: HashMap<Uuid, Request>,
58
59
60
61
    prefill_costs: HashMap<Uuid, PrefillCost>,
    max_num_batched_tokens: Option<usize>,
    active_tokens: usize,
    waiting_tokens: usize,
62
63
64
}

impl SchedulerState {
65
66
67
68
69
70
71
    fn new(max_num_batched_tokens: Option<usize>) -> Self {
        SchedulerState {
            max_num_batched_tokens,
            ..Default::default()
        }
    }

72
73
74
75
    fn is_empty(&self) -> bool {
        self.requests.is_empty()
    }

76
77
78
79
80
81
82
83
84
85
86
    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
    fn receive(&mut self, request: DirectRequest) -> Uuid {
        // Use the provided UUID if available, otherwise generate a new one
        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
        self.requests.insert(uuid, Request::Direct(request));
        self.waiting.push_back(uuid);
        uuid
    }

    /// Get the next UUID from ready or waiting queue and its associated Request.
    fn next(&mut self) -> Option<(Uuid, Request)> {
87
88
89
90
91
        let uuid = self.waiting.pop_front()?;
        let request = self
            .requests
            .remove(&uuid)
            .expect("Request does not exist.");
92
93
94
        Some((uuid, request))
    }

95
96
97
98
99
100
    /// Move a UUID and its Request to the waiting queue (front).
    fn first_in_line(&mut self, uuid: Uuid, request: Request) {
        self.requests.insert(uuid, request);
        self.waiting.push_front(uuid);
    }

101
    /// Move a UUID and its Request to the ready queue.
102
103
    fn move_to_prefill(&mut self, uuid: Uuid, active_seq: ActiveSequence, cost: PrefillCost) {
        self.waiting_tokens += cost.new_tokens;
104
        self.requests.insert(uuid, Request::Active(active_seq));
105
106
        self.prefill.push_back(uuid);
        self.prefill_costs.insert(uuid, cost);
107
108
    }

109
110
111
112
113
    /// Try (chunked) prefill and move to decode queue
    ///
    /// Returns `Some((prefill_compute, creation_signal, is_full_prefill))` where:
    /// - `prefill_compute`: The compute time in milliseconds for this prefill operation
    /// - `creation_signal`: Optional MoveBlock signal for KV cache block creation
114
    /// - `block_hashes`: Block hashes of the sequence beign prefilled
115
    /// - `is_full_prefill`: true if the entire sequence was prefilled, false if chunked
116
    fn try_prefill(&mut self) -> Option<(f64, Option<MoveBlock>, Vec<BlockHash>, bool)> {
117
118
119
        let uuid = self.prefill.pop_front()?;

        // Remove and extract prefill_compute from prefill_costs
120
        let mut prefill_cost = self
121
122
123
            .prefill_costs
            .remove(&uuid)
            .expect("Expects valid prefill cost.");
124

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        let new_tokens = prefill_cost.new_tokens;

        let maybe_prefill_tokens = self.max_num_batched_tokens.and_then(|max_tokens| {
            let remaining_tokens = max_tokens - self.active_tokens;
            if prefill_cost.new_tokens > remaining_tokens {
                Some(remaining_tokens)
            } else {
                None
            }
        });

        let (prefill_compute, is_full_prefill) = if let Some(prefill_tokens) = maybe_prefill_tokens
        {
            let prefill_compute = prefill_cost.predict_prefill_compute(Some(prefill_tokens));
            prefill_cost.new_tokens -= prefill_tokens;
            assert!(
                (prefill_cost.new_tokens > 0) && (prefill_compute > 0.0),
                "Encountered negative prefill tokens or prefill compute cost."
            );

            self.prefill.push_front(uuid);
            self.prefill_costs.insert(uuid, prefill_cost);

            self.active_tokens = self.max_num_batched_tokens.unwrap();
            self.waiting_tokens -= prefill_tokens;

            (prefill_compute, false)
        } else {
            // Assume possible to complete prefilling the sequence, transfer to decode
            self.decode.insert(uuid);

            self.active_tokens += new_tokens;
            self.waiting_tokens -= new_tokens;

            (prefill_cost.predict_prefill_compute(None), true)
        };

        // NOTE: the current behavior allocates the KV blocks for the entire sequence,
        // even if only a chunk is prefilled
        let Some(Request::Active(sequence)) = self.requests.get_mut(&uuid) else {
165
            panic!("Request does not exist.");
166
167
        };

168
169
170
        Some((
            prefill_compute,
            sequence.take_creation_signal(),
171
            sequence.block_hashes(),
172
173
174
175
176
177
178
            is_full_prefill,
        ))
    }

    // assume (chunked) prefills are completed, then active tokens would be 1 per decoding sequence
    fn reset_active_tokens(&mut self) {
        self.active_tokens = self.decode.len();
179
180
    }

181
182
183
184
185
186
187
188
    fn run(&mut self, uuid: Uuid) -> Option<&mut ActiveSequence> {
        if !self.decode.contains(&uuid) {
            return None;
        }
        let Some(Request::Active(sequence)) = self.requests.get_mut(&uuid) else {
            panic!("Request does not exist.");
        };
        Some(sequence)
189
190
    }

191
192
    fn num_active_requests(&self) -> usize {
        self.prefill.len() + self.decode.len()
193
194
195
196
    }

    /// Remove a UUID and its associated Request from collections.
    fn complete(&mut self, uuid: &Uuid) {
197
        tracing::trace!("Request {uuid} will complete");
198
        self.decode.remove(uuid);
199
200
        self.requests.remove(uuid);
        self.prefill_costs.remove(uuid);
201
        self.active_tokens -= 1;
202
203
204
205
206
    }

    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
    /// and adding it back to the waiting queue.
    /// Returns the signal from reset_with_signal or None if no requests are running.
207
    fn preempt(&mut self) -> Vec<MoveBlock> {
208
        // Evict the oldest UUID from running
209
210
211
212
213
214
215
216
        let uuid = self
            .decode
            .evict()
            .expect("Nothing to evict for preemption.");
        let request = self
            .requests
            .remove(&uuid)
            .expect("Request does not exist.");
217
        self.prefill_costs.remove(&uuid);
218
219
        self.active_tokens -= 1;
        tracing::warn!("Request {uuid} will be preempted");
220

221
222
        // Reset the sequence and get the new sequence and signal
        // Insert the new sequence back into the requests map and add to waiting queue
223
224
225
226
227
        let Request::Active(mut active_sequence) = request else {
            panic!("Expected ActiveSequence in running queue")
        };
        let signals = active_sequence.reset_with_signal();

228
229
        // Note: For preemption, we don't compute hit rate since we don't have access to new_tokens
        // and the sequence is being reset anyway. Hit rate tracking is primarily for new scheduling attempts.
230

231
232
233
        self.first_in_line(uuid, Request::Active(active_sequence));

        signals
234
235
236
237
238
239
    }
}

/// Manages scheduling of requests using KvManager resources
#[derive(Clone)]
pub struct Scheduler {
240
    request_tx: mpsc::UnboundedSender<DirectRequest>,
241
    metrics_rx: tokio::sync::watch::Receiver<ForwardPassMetrics>,
242
243
244
245
246
}

impl Scheduler {
    /// Create a new Scheduler with the given parameters
    pub fn new(
247
        args: MockEngineArgs,
Yan Ru Pei's avatar
Yan Ru Pei committed
248
        dp_rank: u32,
249
250
        output_tx: Option<mpsc::UnboundedSender<OutputSignal>>,
        kv_events_tx: Option<mpsc::UnboundedSender<KvCacheEventData>>,
251
252
        cancellation_token: Option<CancellationToken>,
    ) -> Self {
253
254
255
256
257
258
259
        // Create internal channel for KV events only if needed
        let (block_resp_tx, mut block_resp_rx) = if kv_events_tx.is_some() {
            let (tx, rx) = mpsc::unbounded_channel::<MoveBlockResponse>();
            (Some(tx), Some(rx))
        } else {
            (None, None)
        };
260

261
262
263
264
265
266
        // Assert speedup_ratio is greater than 0
        assert!(
            args.speedup_ratio > 0.0,
            "speedup_ratio must be greater than 0, got: {}",
            args.speedup_ratio
        );
267

268
269
        // Create channel for request handling
        let (request_tx, mut request_rx) = mpsc::unbounded_channel::<DirectRequest>();
270
        let mut initial_metrics = ForwardPassMetrics::default();
Yan Ru Pei's avatar
Yan Ru Pei committed
271
        initial_metrics.worker_stats.data_parallel_rank = Some(dp_rank);
272
273
        let (metrics_tx, metrics_rx) =
            tokio::sync::watch::channel::<ForwardPassMetrics>(initial_metrics);
274

275
        let cancel_token_clone = cancellation_token.unwrap_or_default().clone();
276
277
278

        // Spawn main background task with cancellation token
        tokio::spawn(async move {
279
280
281
282
283
            // Create state and kv_manager as local variables owned by this task
            let mut state = SchedulerState::new(args.max_num_batched_tokens);
            let mut kv_manager =
                KvManager::new_with_sender(args.num_gpu_blocks, args.block_size, block_resp_tx);
            let mut hit_rates = VecDeque::with_capacity(1000);
284
            let mut should_schedule = true;
285
286

            loop {
287
288
                {
                    // Enqueue new request, blocks until at least one is received, so no redundant work is done
289
                    if state.is_empty() {
290
291
292
293
                        let Some(request) = request_rx.recv().await else {
                            tracing::warn!("request sender is dropped");
                            break;
                        };
294
                        state.receive(request);
295
296
297
                    }
                }

298
299
300
301
302
303
304
305
                tokio::select! {
                    biased;

                    // Enqueue new request
                    Some(request) = request_rx.recv() => {
                        state.receive(request);
                    }

306
                    // Try Scheduling Requests - runs on normal interval or after simulation
307
                    _ = tokio::task::yield_now() => {
308
309
310
311
312
                        // Skip if we just ran scheduling after simulation to prevent consecutive runs
                        if !should_schedule {
                            continue;
                        }

313
314
                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
                        // schedule anymore.
315
316
317
                        let mut current_blocks = kv_manager.num_active_blocks();
                        let mut current_tokens = state.active_tokens + state.waiting_tokens;
                        let mut current_seqs = state.num_active_requests();
318

319
                        while let Some((uuid, request)) = state.next() {
320
                            let active_sequence = get_active_sequence(request, args.block_size, args.enable_prefix_caching);
321

322
                            // Update predictive budgets
323
                            let prefill_cost = kv_manager.get_prefill_cost(&active_sequence);
324
                            let total_tokens = active_sequence.len();
325
326
                            // this is conservative, assumes no cache hit so never over-schedules
                            let new_blocks = (total_tokens as u32).div_ceil(args.block_size as u32) as usize;
327
328
329
330
331
                            let new_tokens = prefill_cost.new_tokens;

                            current_blocks += new_blocks;
                            current_tokens += new_tokens;
                            current_seqs += 1;
332

333
                            // Check various budgets to see if possible to schedule
334
                            let under_block_budget = current_blocks as f64 <= (1. - args.watermark) * kv_manager.max_capacity() as f64;
335
336
337
                            // If chunked prefill is enabled, we can be under token budget when scheduling
                            let comparison_tokens = if args.enable_chunked_prefill {current_tokens - new_tokens} else {current_tokens};
                            let under_token_budget = args.max_num_batched_tokens.is_none_or(|limit| comparison_tokens <= limit);
338
339
340
341
                            let under_seq_budget = args.max_num_seqs.is_none_or(|limit| current_seqs <= limit);

                            // Cannot schedule, put first in line instead
                            if !(under_block_budget && under_token_budget && under_seq_budget) {
342
                                state.first_in_line(uuid, Request::Active(active_sequence));
343
                                break;
344
345
346
347
                            }

                            // Compute and store hit rate
                            let hit_rate = if !active_sequence.is_empty() { 1.0 - (new_tokens as f32 / active_sequence.len() as f32) } else { 0.0 };
348
349
350
                            hit_rates.push_back(hit_rate);
                            if hit_rates.len() > 1000 {
                                hit_rates.pop_front();
351
                            }
352

353
                            state.move_to_prefill(uuid, active_sequence, prefill_cost);
354
                            should_schedule = false;
355
356
357
358
                        }
                    }

                    // Check for cancellation
359
                    _ = cancel_token_clone.cancelled() => {
360
361
                        break;
                    }
362
                }
363

364
365
                // Simulates prefill + decode
                // Base time needed for decoding using active percentage and quadratic formula
366
                let active_perc = kv_manager.get_active_perc();
367
368
369
370
                let decoding_time = -5.47 * active_perc.powi(2) + 43.88 * active_perc + 19.44;
                let mut total_time = Duration::from_secs_f64(decoding_time / 1000.0);

                // Process prefilling
371
372
373
374
375
                while let Some((
                    prefill_compute,
                    maybe_creation_signal,
                    block_hashes,
                    is_full_prefill,
376
                )) = state.try_prefill()
377
378
379
380
381
382
                {
                    // NOTE: Prefill cost/time is always incremented for new blocks, even if they
                    // could be cached by other requests in the same batch. This matches vLLM behavior.
                    total_time += Duration::from_secs_f64(prefill_compute / 1000.0);

                    if let Some(creation_signal) = maybe_creation_signal {
383
384
                        if !process_signals(&mut kv_manager, std::slice::from_ref(&creation_signal))
                        {
385
                            panic!("Block allocation for prefilling cannot fail.");
386
                        }
387

388
                        // Drain KV events and forward to relay after prefill signal processing
389
                        if let (Some(relay_tx), Some(rx)) = (&kv_events_tx, &mut block_resp_rx) {
390
                            while let Ok(event) = rx.try_recv() {
391
392
                                let _ =
                                    relay_tx.send(block_response_to_kv_event(event, &block_hashes));
393
                            }
394
395
                        }
                    };
396

397
398
399
400
401
                    // Impossible to schedule more prefills if we encounter one incomplete (chunked) prefill
                    if !is_full_prefill {
                        break;
                    }
                }
402

403
                state.reset_active_tokens();
404

405
                // Process decoding
406
                let uuids: Vec<Uuid> = state.decode.keys().cloned().collect();
407
408
409
410
                if !uuids.is_empty() {
                    should_schedule = true
                };
                for uuid in uuids {
411
                    let Some(sequence) = state.run(uuid) else {
412
413
414
415
416
417
                        continue;
                    };
                    let signals = sequence.generate();

                    // Process all signals with the KvManager
                    // Handling of preemption on failure
418
                    if !process_signals(&mut kv_manager, &signals) {
419
                        sequence.pop(); // revert the failed generation op
420
421
                        for signal in state.preempt() {
                            kv_manager.process(&signal);
422
423
424
                        }
                        continue;
                    }
425

426
                    // Drain KV events and forward to relay after decode signal processing
427
                    if let (Some(relay_tx), Some(rx)) = (&kv_events_tx, &mut block_resp_rx) {
428
                        while let Ok(event) = rx.try_recv() {
429
430
                            let _ = relay_tx
                                .send(block_response_to_kv_event(event, &sequence.block_hashes()));
431
432
                        }
                    }
433

434
435
436
437
438
439
440
                    // Check completion and send notification
                    let is_complete = sequence.generated_tokens() >= sequence.max_output_tokens();
                    let should_output =
                        sequence.generated_tokens() > sequence.already_generated_tokens();

                    let mut send_failed = false;
                    if should_output {
441
                        send_failed = output_tx.as_ref().is_some_and(|tx| {
442
443
444
445
446
447
448
                            tx.send(OutputSignal {
                                uuid,
                                completed: is_complete,
                            })
                            .is_err()
                        });
                    }
449

450
451
                    if send_failed {
                        for signal in &sequence.free_signal() {
452
                            kv_manager.process(signal);
453
                        }
454
                    }
455

456
                    if send_failed || is_complete {
457
                        state.complete(&uuid);
458
                        continue;
459
460
                    }
                }
461

462
463
464
465
466
467
                // Send metrics once per forward pass (after all prefill and decode processing)
                {
                    let metrics = get_fwd_pass_metrics(&state, &kv_manager, &hit_rates, dp_rank);
                    let _ = metrics_tx.send(metrics);
                }

468
469
470
471
472
473
                // Sleep once for the adjusted duration
                let adjusted_time =
                    Duration::from_secs_f64(total_time.as_secs_f64() / args.speedup_ratio);
                if adjusted_time.as_millis() > 0 {
                    tokio::time::sleep(adjusted_time).await;
                }
474
475
476
477
478
            }
        });

        Self {
            request_tx,
479
            metrics_rx,
480
481
482
483
484
        }
    }

    /// Add a new request to the waiting queue
    pub async fn receive(&self, request: DirectRequest) {
485
486
487
488
489
        let _ = self.request_tx.send(request);
    }

    pub fn request_sender(&self) -> mpsc::UnboundedSender<DirectRequest> {
        self.request_tx.clone()
490
491
    }

492
493
494
495
496
    /// Get a watch receiver for forward pass metrics
    pub fn metrics_receiver(&self) -> tokio::sync::watch::Receiver<ForwardPassMetrics> {
        self.metrics_rx.clone()
    }
}
497

498
499
500
501
502
/// Calculate forward pass metrics from current state
fn get_fwd_pass_metrics(
    state: &SchedulerState,
    kv_manager: &KvManager,
    hit_rates: &VecDeque<f32>,
Yan Ru Pei's avatar
Yan Ru Pei committed
503
    dp_rank: u32,
504
505
506
507
508
509
510
511
512
513
514
515
516
) -> ForwardPassMetrics {
    // Get state metrics
    let request_active_slots = state.decode.len() as u64;
    let num_requests_waiting = state.waiting.len() as u64;

    // Get KV manager metrics
    let active_blocks_count = kv_manager.active_blocks().len() as u64;
    let total_capacity = kv_manager.max_capacity() as u64;
    let gpu_cache_usage_perc = if total_capacity > 0 {
        active_blocks_count as f32 / total_capacity as f32
    } else {
        0.0
    };
517

518
519
520
521
522
523
524
    // Get hit rate metrics
    let gpu_prefix_cache_hit_rate = if hit_rates.is_empty() {
        0.0
    } else {
        let sum: f32 = hit_rates.iter().sum();
        sum / hit_rates.len() as f32
    };
525

526
    let worker_stats = WorkerStats {
Yan Ru Pei's avatar
Yan Ru Pei committed
527
        data_parallel_rank: Some(dp_rank),
528
529
530
531
        request_active_slots,
        request_total_slots: 1024, // vllm max_num_seqs for gpu >= 70 vram, otherwise 256, fallback is 128
        num_requests_waiting,
    };
532

533
534
535
536
537
538
    let kv_stats = KvStats {
        kv_active_blocks: active_blocks_count,
        kv_total_blocks: total_capacity,
        gpu_cache_usage_perc,
        gpu_prefix_cache_hit_rate,
    };
539

540
    let spec_decode_stats = None;
541

542
543
544
545
    ForwardPassMetrics {
        worker_stats,
        kv_stats,
        spec_decode_stats,
546
547
548
549
    }
}

/// Convert a Request to an ActiveSequence
550
551
552
553
554
fn get_active_sequence(
    request: Request,
    block_size: usize,
    enable_prefix_caching: bool,
) -> ActiveSequence {
555
556
557
558
559
560
561
562
563
564
565
566
    if let Request::Active(active_seq) = request {
        return active_seq;
    }

    let Request::Direct(direct_request) = request else {
        unreachable!("Request must be either Direct or Active");
    };

    ActiveSequence::new(
        direct_request.tokens,
        direct_request.max_output_tokens,
        Some(block_size),
567
        enable_prefix_caching,
568
569
570
571
572
573
574
575
576
577
    )
}

/// Processes MoveBlock signals with the KvManager.
///
/// When a signal fails, this function verifies that the failure is for an expected case:
/// specifically a single signal attempting to create a single partial (generation) block.
/// This validation is important because in normal operation, the only legitimate failure
/// case should be when trying to acquire a new generation block - any other failures would
/// indicate an unexpected state in the system.
578
fn process_signals(kv_manager: &mut KvManager, signals: &[MoveBlock]) -> bool {
579
    for signal in signals {
580
        if kv_manager.process(signal) {
581
582
583
584
            continue;
        }

        // Check we have a Use signal with blocks
585
        let MoveBlock::Use(blocks) = signal else {
586
587
588
            panic!(
                "Failed signal is Invalid. Has to fail on generation signal, but failed on {signal:?}"
            );
589
590
591
        };

        // Verify the signal contains exactly one block
592
        let num_blocks = blocks.len();
593
        let num_active_blocks = kv_manager.num_active_blocks();
594
595
596
597
        if num_blocks != 1 {
            panic!(
                "Failed signal is Invalid. Tried to create (prefill) {num_blocks} blocks on top of {num_active_blocks} active blocks."
            );
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
        }

        // Verify the block is a PartialBlock (generation block)
        if !matches!(blocks[0], UniqueBlock::PartialBlock(_)) {
            panic!("Failed signal is Invalid. Generation block has to be partial.");
        }

        return false;
    }

    true
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;
    use std::time::Duration;
616
    use tokio::time::interval;
617

618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
    /// Helper function to verify that the scheduler is idle (no active or waiting requests/resources)
    fn assert_scheduler_idle(metrics: &ForwardPassMetrics) {
        assert_eq!(
            metrics.worker_stats.request_active_slots, 0,
            "Expected 0 active slots, got {}",
            metrics.worker_stats.request_active_slots
        );
        assert_eq!(
            metrics.worker_stats.num_requests_waiting, 0,
            "Expected 0 waiting requests, got {}",
            metrics.worker_stats.num_requests_waiting
        );
        assert_eq!(
            metrics.kv_stats.kv_active_blocks, 0,
            "Expected 0 active blocks, got {}",
            metrics.kv_stats.kv_active_blocks
        );
        assert_eq!(
            metrics.kv_stats.gpu_cache_usage_perc, 0.0,
            "Expected 0% GPU cache usage, got {}",
            metrics.kv_stats.gpu_cache_usage_perc
        );
    }

642
    #[rstest]
643
644
645
646
647
648
649
650
    #[case::case_1(false, false, false)]
    #[case::case_2(false, true, false)]
    #[case::case_3(true, false, false)]
    #[case::case_4(true, true, false)]
    #[case::case_5(false, false, true)]
    #[case::case_6(false, true, true)]
    #[case::case_7(true, false, true)]
    #[case::case_8(true, true, true)]
651
    #[tokio::test]
652
653
654
    async fn test_scheduler_token_generation_patterns(
        #[case] use_shared_tokens: bool,
        #[case] enable_prefix_caching: bool,
655
        #[case] enable_chunked_prefill: bool,
656
    ) {
657
        unsafe { std::env::set_var("RUST_LOG", "debug") };
658
659

        let kv_capacity: usize = 500;
660
        let block_size: usize = 64;
661
        let num_requests: usize = 200;
662
663
664
665
        let input_len: usize = 1000;
        let max_output_tokens: usize = 100;

        // Create channel for token output
666
667
668
669
670
671
672
673
        let (output_tx, mut output_rx) = mpsc::unbounded_channel::<OutputSignal>();

        // Create scheduler args using builder - now including enable_prefix_caching
        let args = MockEngineArgs::builder()
            .num_gpu_blocks(kv_capacity)
            .block_size(block_size)
            .speedup_ratio(10.0)
            .enable_prefix_caching(enable_prefix_caching)
674
            .enable_chunked_prefill(enable_chunked_prefill)
675
676
677
678
            .build()
            .unwrap();

        // Create scheduler with new args struct
Yan Ru Pei's avatar
Yan Ru Pei committed
679
        let scheduler = Scheduler::new(args, 0, Some(output_tx), None, None);
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709

        // Create shared tokens for caching case
        let shared_tokens = if use_shared_tokens {
            Some(
                (0..input_len / 2)
                    .map(|_| rand::random::<u32>() % 50000)
                    .collect::<Vec<_>>(),
            )
        } else {
            None
        };

        // Create test requests
        for _ in 0..num_requests {
            let input_tokens = if let Some(ref shared) = shared_tokens {
                // For caching case: use shared tokens for first half, random for second half
                let mut tokens = shared.clone();
                tokens.extend((0..input_len / 2).map(|_| rand::random::<u32>() % 50000));
                tokens
            } else {
                // For random case: create unique random token vector for each request
                (0..input_len)
                    .map(|_| rand::random::<u32>() % 50000)
                    .collect::<Vec<_>>()
            };

            let request = DirectRequest {
                tokens: input_tokens,
                max_output_tokens,
                uuid: None,
Yan Ru Pei's avatar
Yan Ru Pei committed
710
                dp_rank: 0,
711
712
713
714
715
716
717
718
719
720
721
722
723
724
            };
            scheduler.receive(request).await;
        }

        let start_time = std::time::Instant::now();

        // Collect all generated tokens (should be num_requests * max_output_tokens)
        let expected_tokens = num_requests * max_output_tokens;
        let mut received_tokens = 0;

        // Set up a timeout that causes the test to panic if no tokens are received for 2 seconds
        let timeout = tokio::time::sleep(Duration::from_secs(2));
        tokio::pin!(timeout);

725
726
727
        // Get metrics receiver
        let metrics_rx = scheduler.metrics_receiver();

728
729
730
731
732
733
734
735
736
        // Set up debug ticker interval
        let mut debug_interval = interval(Duration::from_millis(500));

        loop {
            tokio::select! {
                biased;

                // Manual debug ticker that prints forward pass metrics
                _ = debug_interval.tick() => {
737
                    let _metrics = metrics_rx.borrow().clone();
738
                    tracing::debug!("Forward Pass Metrics: {_metrics:#?}");
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
                }

                Some(_) = output_rx.recv() => {
                    received_tokens += 1;
                    // Reset timeout whenever we receive a token
                    timeout.set(tokio::time::sleep(Duration::from_secs(2)));
                }

                _ = &mut timeout => {
                    // Break instead of panicking when timeout occurs
                    break;
                }
            }
        }

        // Calculate and print elapsed time
        let elapsed = start_time.elapsed();
        println!(
757
            "Test completed in: {elapsed:?} for {} case with prefix_caching={enable_prefix_caching} and chunked_prefill={enable_chunked_prefill}",
758
759
760
761
            if use_shared_tokens {
                "caching"
            } else {
                "random"
762
            }
763
764
765
766
        );

        // Assert that we received the expected number of tokens
        assert!(
767
768
769
            received_tokens == expected_tokens,
            "Received {received_tokens} tokens but expected exactly {expected_tokens}"
        );
770

771
772
        // Wait a bit for final metrics update to propagate
        tokio::time::sleep(Duration::from_millis(100)).await;
773

774
775
        let metrics = scheduler.metrics_receiver().borrow().clone();
        assert_scheduler_idle(&metrics);
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
    }

    #[tokio::test]
    async fn test_cache_hit_rate_with_identical_requests() {
        let block_size: usize = 64;
        let max_output_tokens: usize = 10;
        let speedup_ratio = 10.0;
        let num_requests = 10;
        let token_length = 65;

        // Create channel for token output
        let (output_tx, mut output_rx) = mpsc::unbounded_channel::<OutputSignal>();

        // Create scheduler args
        let args = MockEngineArgs::builder()
            .num_gpu_blocks(100) // Large enough to not be a constraint
            .block_size(block_size)
            .speedup_ratio(speedup_ratio)
            .build()
            .unwrap();

        // Create scheduler
Yan Ru Pei's avatar
Yan Ru Pei committed
798
        let scheduler = Scheduler::new(args, 0, Some(output_tx), None, None);
799
800
801
802
803
804
805
806
807
808

        // Create identical tokens for all requests
        let identical_tokens: Vec<u32> = (0..token_length).map(|i| i as u32).collect();

        // Send all requests with identical tokens
        for _ in 0..num_requests {
            let request = DirectRequest {
                tokens: identical_tokens.clone(),
                max_output_tokens,
                uuid: None,
Yan Ru Pei's avatar
Yan Ru Pei committed
809
                dp_rank: 0,
810
811
812
813
814
815
816
817
818
819
820
821
822
            };
            scheduler.receive(request).await;
            // Sleep for 0.1 second after each request
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        // Collect all generated tokens
        let mut received_tokens = 0;

        // Set up a timeout that resets to 0.5 seconds on each received token
        let timeout = tokio::time::sleep(Duration::from_millis(500));
        tokio::pin!(timeout);

823
824
825
        // Get metrics receiver
        let metrics_rx = scheduler.metrics_receiver();

826
827
828
829
830
831
832
833
834
        // Set up debug ticker interval
        let mut debug_interval = interval(Duration::from_millis(500));

        loop {
            tokio::select! {
                biased;

                // Manual debug ticker that prints forward pass metrics
                _ = debug_interval.tick() => {
835
                    let _metrics = metrics_rx.borrow().clone();
836
                    tracing::debug!("Forward Pass Metrics: {_metrics:#?}");
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
                }

                Some(_signal) = output_rx.recv() => {
                    received_tokens += 1;
                    // Reset timeout whenever we receive a token
                    timeout.set(tokio::time::sleep(Duration::from_millis(500)));
                }

                _ = &mut timeout => {
                    // Break when timeout occurs (no more tokens for 0.5 seconds)
                    break;
                }
            }
        }

852
853
854
        // Wait a bit for final metrics update
        tokio::time::sleep(Duration::from_millis(100)).await;

855
        // Verify forward pass metrics
856
        let metrics = metrics_rx.borrow().clone();
857

858
        assert_scheduler_idle(&metrics);
859
        assert!(
860
            metrics.kv_stats.gpu_prefix_cache_hit_rate > 0.8,
861
            "Expected cache hit rate > 0.8, got {}",
862
            metrics.kv_stats.gpu_prefix_cache_hit_rate
863
864
865
866
        );

        println!(
            "Test passed! Cache hit rate: {:.3}",
867
            metrics.kv_stats.gpu_prefix_cache_hit_rate
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
        );
        println!("Received {received_tokens} tokens");
    }

    #[tokio::test]
    async fn test_receiver_drop_cleans_up_resources() {
        let block_size: usize = 64;
        let input_tokens = 256;
        let max_output_tokens = 200; // More than we'll receive

        // Create channel for token output
        let (output_tx, mut output_rx) = mpsc::unbounded_channel::<OutputSignal>();

        // Create scheduler args
        let args = MockEngineArgs::builder()
            .num_gpu_blocks(10) // Enough for 256 tokens (4 blocks)
            .block_size(block_size)
            .speedup_ratio(100.0) // Fast simulation
            .build()
            .unwrap();

        // Create scheduler
Yan Ru Pei's avatar
Yan Ru Pei committed
890
        let scheduler = Scheduler::new(args, 0, Some(output_tx), None, None);
891
892
893
894
895
896
897

        // Create request with 256 tokens
        let tokens: Vec<u32> = (0..input_tokens).map(|i| i as u32).collect();
        let request = DirectRequest {
            tokens,
            max_output_tokens,
            uuid: None,
Yan Ru Pei's avatar
Yan Ru Pei committed
898
            dp_rank: 0,
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
        };

        scheduler.receive(request).await;

        // Receive exactly 129 tokens
        let mut received_count = 0;
        while received_count < 129 {
            if let Some(_signal) = output_rx.recv().await {
                received_count += 1;
            } else {
                panic!("Channel closed before receiving 129 tokens");
            }
        }

        // Drop the receiver immediately
        drop(output_rx);

        // Wait for 1 second to allow cleanup
        tokio::time::sleep(Duration::from_secs(1)).await;

        // Check forward pass metrics
920
921
        let metrics_rx = scheduler.metrics_receiver();
        let metrics = metrics_rx.borrow().clone();
922

923
        assert_scheduler_idle(&metrics);
924
925
    }
}