scheduler.rs 32 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
31
32
33
34
35
36
37
38
39
40
41
42
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! 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

43
use crate::kv_router::protocols::{ForwardPassMetrics, KvCacheEventData};
44
45
use crate::mocker::evictor::LRUEvictor;
use crate::mocker::kv_manager::KvManager;
46
47
48
49
use crate::mocker::protocols::{
    block_response_to_kv_event, MoveBlock, OutputSignal, PrefillCost, UniqueBlock,
};
use crate::mocker::protocols::{DirectRequest, MockEngineArgs, MoveBlockResponse};
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
use crate::mocker::sequence::ActiveSequence;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use tokio::time::{interval, Duration};
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>,
68
69
    prefill: VecDeque<Uuid>,
    decode: LRUEvictor<Uuid>,
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
    requests: HashMap<Uuid, Request>,
    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
}

impl SchedulerState {
    /// 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)> {
86
87
88
89
90
        let uuid = self.waiting.pop_front()?;
        let request = self
            .requests
            .remove(&uuid)
            .expect("Request does not exist.");
91
92
93
        Some((uuid, request))
    }

94
95
96
97
98
99
    /// 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);
    }

100
    /// Move a UUID and its Request to the ready queue.
101
    fn start_prefill(&mut self, uuid: Uuid, active_seq: ActiveSequence, cost: Option<PrefillCost>) {
102
        self.requests.insert(uuid, Request::Active(active_seq));
103
104
        self.prefill.push_back(uuid);
        self.prefill_costs.insert(uuid, cost);
105
106
    }

107
108
109
110
111
112
113
114
115
116
117
118
    /// Pop from prefill queue and move to decode queue.
    /// Returns the prefill_compute value if available.
    fn start_decode(&mut self) -> Option<(f64, MoveBlock)> {
        let uuid = self.prefill.pop_front()?;
        self.decode.insert(uuid);

        // Remove and extract prefill_compute from prefill_costs
        let prefill_cost = self
            .prefill_costs
            .remove(&uuid)
            .flatten()
            .expect("Expects valid prefill cost.");
119
120

        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
121
            panic!("Request does not exist.");
122
        };
123
124
125
126
        let creation_signal = sequence
            .creation_signal()
            .clone()
            .expect("Must have creation signal.");
127

128
        Some((prefill_cost.prefill_compute, creation_signal))
129
130
    }

131
132
133
134
135
136
137
138
    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)
139
140
    }

141
142
    fn num_active_requests(&self) -> usize {
        self.prefill.len() + self.decode.len()
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    }

    /// Calculate the current running batched tokens
    fn num_batched_tokens(&self) -> usize {
        self.prefill_costs
            .values()
            .map(|cost| match cost {
                Some(cost) => cost.new_tokens,
                None => 1,
            })
            .sum()
    }

    /// Remove a UUID and its associated Request from collections.
    fn complete(&mut self, uuid: &Uuid) {
        // println!("Request {} will complete", uuid);
159
        self.decode.remove(uuid);
160
161
162
163
164
165
166
        self.requests.remove(uuid);
        self.prefill_costs.remove(uuid);
    }

    /// 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.
167
    fn preempt(&mut self) -> Vec<MoveBlock> {
168
        // Evict the oldest UUID from running
169
170
171
172
173
174
175
176
        let uuid = self
            .decode
            .evict()
            .expect("Nothing to evict for preemption.");
        let request = self
            .requests
            .remove(&uuid)
            .expect("Request does not exist.");
177
        self.prefill_costs.remove(&uuid);
178
        eprintln!("Request {uuid} will be preempted");
179

180
181
        // Reset the sequence and get the new sequence and signal
        // Insert the new sequence back into the requests map and add to waiting queue
182
183
184
185
186
        let Request::Active(mut active_sequence) = request else {
            panic!("Expected ActiveSequence in running queue")
        };
        let signals = active_sequence.reset_with_signal();

187
188
        // 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.
189

190
191
192
        self.first_in_line(uuid, Request::Active(active_sequence));

        signals
193
194
195
196
197
198
    }
}

/// Manages scheduling of requests using KvManager resources
#[derive(Clone)]
pub struct Scheduler {
199
    dp_rank: Option<u32>,
200
201
    state: Arc<Mutex<SchedulerState>>,
    kv_manager: Arc<Mutex<KvManager>>,
202
203
    request_tx: mpsc::UnboundedSender<DirectRequest>,
    hit_rates: Arc<Mutex<VecDeque<f32>>>,
204
205
206
207
208
}

impl Scheduler {
    /// Create a new Scheduler with the given parameters
    pub fn new(
209
210
211
212
        args: MockEngineArgs,
        dp_rank: Option<u32>,
        output_tx: Option<mpsc::UnboundedSender<OutputSignal>>,
        kv_events_tx: Option<mpsc::UnboundedSender<KvCacheEventData>>,
213
214
215
216
        cancellation_token: Option<CancellationToken>,
    ) -> Self {
        let state = Arc::new(Mutex::new(SchedulerState::default()));

217
218
219
220
221
222
223
        // 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)
        };
224

225
226
227
228
229
230
231
232
233
234
235
236
237
        let kv_manager = Arc::new(Mutex::new(KvManager::new_with_sender(
            args.num_gpu_blocks,
            args.block_size,
            block_resp_tx,
        )));
        let hit_rates = Arc::new(Mutex::new(VecDeque::with_capacity(1000)));

        // Assert speedup_ratio is greater than 0
        assert!(
            args.speedup_ratio > 0.0,
            "speedup_ratio must be greater than 0, got: {}",
            args.speedup_ratio
        );
238

239
240
        // Create channel for request handling
        let (request_tx, mut request_rx) = mpsc::unbounded_channel::<DirectRequest>();
241
242
243
244
245

        // Create a clone for the background task
        let state_clone = state.clone();
        let kv_manager_clone = kv_manager.clone();
        let output_tx_clone = output_tx.clone();
246
247
        let cancel_token_clone = cancellation_token.unwrap_or_default().clone();
        let hit_rates_clone = hit_rates.clone();
248
249
250

        // Spawn main background task with cancellation token
        tokio::spawn(async move {
251
252
253
            let mut schedule_interval = interval(Duration::from_secs_f64(1e-3));
            let mut simulate_interval = interval(Duration::from_secs_f64(1e-4));
            let mut should_schedule = true;
254
255
256
257
258
259
260
261
262
263
264

            loop {
                tokio::select! {
                    biased;

                    // Enqueue new request
                    Some(request) = request_rx.recv() => {
                        let mut state = state_clone.lock().await;
                        state.receive(request);
                    }

265
                    // Try Scheduling Requests - runs on normal interval or after simulation
266
                    _ = schedule_interval.tick() => {
267
268
269
270
271
                        // Skip if we just ran scheduling after simulation to prevent consecutive runs
                        if !should_schedule {
                            continue;
                        }

272
                        let mut state_guard = state_clone.lock().await;
273
                        let kv_manager_guard = kv_manager_clone.lock().await;
274
275
276

                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
                        // schedule anymore.
277
278
279
280
                        let mut current_blocks = kv_manager_guard.num_active_blocks();
                        let mut current_tokens = state_guard.num_batched_tokens();
                        let mut current_seqs = state_guard.num_active_requests();

281
                        while let Some((uuid, request)) = state_guard.next() {
282
                            let active_sequence = get_active_sequence(request, args.block_size, args.enable_prefix_caching);
283

284
285
286
287
288
289
290
291
292
                            // Update predictive budgets
                            let prefill_cost = kv_manager_guard.get_prefill_cost(&active_sequence);
                            let total_tokens = active_sequence.len();
                            let new_blocks = (total_tokens + 1) / args.block_size;  // this is conservative, assumes no cache hit
                            let new_tokens = prefill_cost.new_tokens;

                            current_blocks += new_blocks;
                            current_tokens += new_tokens;
                            current_seqs += 1;
293
294

                            // Check if it can be scheduled
295
296
297
298
299
300
301
                            let under_block_budget = current_blocks as f64 <= (1. - args.watermark) * kv_manager_guard.max_capacity() as f64;
                            let under_token_budget = args.max_num_batched_tokens.is_none_or(|limit| current_tokens <= limit);
                            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) {
                                state_guard.first_in_line(uuid, Request::Active(active_sequence));
302
                                break;
303
304
305
306
307
308
309
310
311
312
313
                            }

                            // 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 };
                            {
                                let mut hit_rates_guard = hit_rates_clone.lock().await;
                                hit_rates_guard.push_back(hit_rate);
                                if hit_rates_guard.len() > 1000 {
                                    hit_rates_guard.pop_front();
                                }
                            }
314

315
316
                            state_guard.start_prefill(uuid, active_sequence, Some(prefill_cost));
                            should_schedule = false;
317
318
319
320
                        }
                    }

                    // Check for cancellation
321
                    _ = cancel_token_clone.cancelled() => {
322
323
324
325
326
327
328
329
                        break;
                    }

                    // Simulate running requests (prefill + decode)
                    _ = simulate_interval.tick() => {
                        let mut state_guard = state_clone.lock().await;
                        let mut kv_manager_guard = kv_manager_clone.lock().await;

330
331
332
333
334
335
336
337
338
339
340
341
342
                        // Base time needed for decoding using active percentage and quadratic formula
                        let active_perc = kv_manager_guard.get_active_perc();
                        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
                        while let Some((prefill_compute, creation_signal)) = state_guard.start_decode() {
                            // 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);
                            let prefill_success = process_signals(&mut kv_manager_guard, std::slice::from_ref(&creation_signal));
                            if !prefill_success {
                                panic!("Block allocation for prefilling cannot fail.");
343
344
                            }

345
346
347
348
349
350
351
                            // Drain KV events and forward to relay after prefill signal processing
                            if let (Some(ref relay_tx), Some(ref mut rx)) = (&kv_events_tx, &mut block_resp_rx) {
                                while let Ok(event) = rx.try_recv() {
                                    let _ = relay_tx.send(block_response_to_kv_event(event));
                                }
                            }
                        }
352

353
354
355
356
357
358
359
                        // Process decoding
                        let uuids: Vec<Uuid> = state_guard.decode.keys().cloned().collect();
                        if !uuids.is_empty() {should_schedule = true};
                        for uuid in uuids {
                            let Some(sequence) = state_guard.run(uuid) else {
                                continue;
                            };
360
361
362
363
364
365
                            let signals = sequence.generate();

                            // Process all signals with the KvManager
                            // Handling of preemption on failure
                            if !process_signals(&mut kv_manager_guard, &signals) {
                                sequence.pop();  // revert the failed generation op
366
                                for signal in state_guard.preempt() {
367
368
369
370
371
                                    kv_manager_guard.process(&signal);
                                }
                                continue;
                            }

372
373
374
375
376
                            // Drain KV events and forward to relay after decode signal processing
                            if let (Some(ref relay_tx), Some(ref mut rx)) = (&kv_events_tx, &mut block_resp_rx) {
                                while let Ok(event) = rx.try_recv() {
                                    let _ = relay_tx.send(block_response_to_kv_event(event));
                                }
377
378
                            }

379
380
381
382
383
384
385
386
387
                            // 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 {
                                send_failed = output_tx_clone.as_ref().is_some_and(|tx| {
                                    tx.send(OutputSignal { uuid, completed: is_complete }).is_err()
                                });
388
389
                            }

390
391
392
393
394
395
396
397
398
                            if send_failed {
                                for signal in &sequence.free_signal() {
                                    kv_manager_guard.process(signal);
                                }
                            }

                            if send_failed || is_complete {
                                state_guard.complete(&uuid);
                                continue;
399
400
401
                            }
                        }

402
403
404
405
406
407
                        // Sleep once for the adjusted duration
                        drop(kv_manager_guard);
                        drop(state_guard);
                        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;
408
409
410
411
412
413
414
                        }
                    }
                }
            }
        });

        Self {
415
            dp_rank,
416
417
418
            state,
            kv_manager,
            request_tx,
419
            hit_rates,
420
421
422
423
424
        }
    }

    /// Add a new request to the waiting queue
    pub async fn receive(&self, request: DirectRequest) {
425
426
427
428
429
430
        let _ = self.request_tx.send(request);
    }

    /// Expose the sender
    pub fn request_sender(&self) -> mpsc::UnboundedSender<DirectRequest> {
        self.request_tx.clone()
431
432
433
434
435
436
437
438
439
440
441
    }

    /// Get the count of waiting requests
    pub async fn waiting_count(&self) -> usize {
        let state = self.state.lock().await;
        state.waiting.len()
    }

    /// Get the count of running requests
    pub async fn running_count(&self) -> usize {
        let state = self.state.lock().await;
442
        state.decode.len()
443
444
445
446
447
448
449
450
451
452
    }

    /// Get the current capacity of the KvManager
    pub async fn kv_usage_perc(&self) -> f64 {
        let kv_manager = self.kv_manager.lock().await;
        kv_manager.current_capacity_perc()
    }

    /// Returns forward pass metrics for monitoring purposes
    pub async fn get_forward_pass_metrics(&self) -> ForwardPassMetrics {
453
        // Acquire all locks in consistent order: state -> kv_manager -> hit_rates
454
455
        let state = self.state.lock().await;
        let kv_manager = self.kv_manager.lock().await;
456
457
458
459
460
        let hit_rates_guard = self.hit_rates.lock().await;

        // Get state metrics
        let request_active_slots = state.decode.len() as u64;
        let num_requests_waiting = state.waiting.len() as u64;
461

462
        // Get KV manager metrics
463
464
465
466
467
468
469
470
        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
        };

471
472
473
474
475
476
477
478
        // Get hit rate metrics
        let gpu_prefix_cache_hit_rate = if hit_rates_guard.is_empty() {
            0.0
        } else {
            let sum: f32 = hit_rates_guard.iter().sum();
            sum / hit_rates_guard.len() as f32
        };

479
        ForwardPassMetrics {
480
481
482
483
            data_parallel_rank: self.dp_rank,
            request_active_slots,
            // vllm max_num_seqs for gpu >= 70 vram, otherwise 256, fallback is 128
            request_total_slots: 1024,
484
485
            kv_active_blocks: active_blocks_count,
            kv_total_blocks: total_capacity,
486
            num_requests_waiting,
487
            gpu_cache_usage_perc,
488
            gpu_prefix_cache_hit_rate,
489
        }
490
        // Guards drop naturally here in reverse order (LIFO): hit_rates_guard, kv_manager, state
491
492
493
494
    }
}

/// Convert a Request to an ActiveSequence
495
496
497
498
499
fn get_active_sequence(
    request: Request,
    block_size: usize,
    enable_prefix_caching: bool,
) -> ActiveSequence {
500
501
502
503
504
505
506
507
508
509
510
511
    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),
512
        enable_prefix_caching,
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
    )
}

/// 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.
fn process_signals(
    kv_manager_guard: &mut tokio::sync::MutexGuard<'_, KvManager>,
    signals: &[MoveBlock],
) -> bool {
    for signal in signals {
        if kv_manager_guard.process(signal) {
            continue;
        }

        // Check we have a Use signal with blocks
533
        let MoveBlock::Use(blocks) = signal else {
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
            panic!("Failed signal is Invalid. Has to fail on generation signal.");
        };

        // Verify the signal contains exactly one block
        if blocks.len() != 1 {
            panic!("Failed signal is Invalid. Can have only one generation signal.");
        }

        // 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;

    #[rstest]
560
561
562
563
    #[case::random_no_prefix_caching(false, false)]
    #[case::random_with_prefix_caching(false, true)]
    #[case::caching_no_prefix_caching(true, false)]
    #[case::caching_with_prefix_caching(true, true)]
564
    #[tokio::test]
565
566
567
568
    async fn test_scheduler_token_generation_patterns(
        #[case] use_shared_tokens: bool,
        #[case] enable_prefix_caching: bool,
    ) {
569
570
571
        std::env::set_var("RUST_LOG", "debug");

        let kv_capacity: usize = 500;
572
        let block_size: usize = 64;
573
574
575
576
577
        let num_requests: usize = 100;
        let input_len: usize = 1000;
        let max_output_tokens: usize = 100;

        // Create channel for token output
578
579
580
581
582
583
584
585
586
587
588
589
590
        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)
            .build()
            .unwrap();

        // Create scheduler with new args struct
        let scheduler = Scheduler::new(args, None, Some(output_tx), None, None);
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620

        // 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,
621
                dp_rank: None,
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
            };
            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);

        // 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() => {
                    let _metrics = scheduler.get_forward_pass_metrics().await;
646
                    println!("Forward Pass Metrics: {_metrics:#?}");
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
                }

                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!(
665
            "Test completed in: {:?} for {} case with prefix_caching={}",
666
667
668
669
670
            elapsed,
            if use_shared_tokens {
                "caching"
            } else {
                "random"
671
672
            },
            enable_prefix_caching
673
674
675
676
        );

        // Assert that we received the expected number of tokens
        assert!(
677
678
679
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
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
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
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
            received_tokens == expected_tokens,
            "Received {received_tokens} tokens but expected exactly {expected_tokens}"
        );
    }

    #[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
        let scheduler = Scheduler::new(args, None, Some(output_tx), None, None);

        // 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,
                dp_rank: None,
            };
            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);

        // 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() => {
                    let _metrics = scheduler.get_forward_pass_metrics().await;
                    println!("Forward Pass Metrics: {_metrics:#?}");
                }

                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;
                }
            }
        }

        // Verify forward pass metrics
        let metrics = scheduler.get_forward_pass_metrics().await;

        assert_eq!(
            metrics.num_requests_waiting, 0,
            "Expected no waiting requests, got {}",
            metrics.num_requests_waiting
        );

        assert!(
            metrics.gpu_prefix_cache_hit_rate > 0.8,
            "Expected cache hit rate > 0.8, got {}",
            metrics.gpu_prefix_cache_hit_rate
        );

        println!(
            "Test passed! Cache hit rate: {:.3}",
            metrics.gpu_prefix_cache_hit_rate
        );
        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
        let scheduler = Scheduler::new(args, None, Some(output_tx), None, None);

        // 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,
            dp_rank: None,
        };

        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
        let metrics = scheduler.get_forward_pass_metrics().await;

        assert_eq!(
            metrics.gpu_cache_usage_perc,
            0.0,
            "Expected GPU cache usage to be 0%, got {}%",
            metrics.gpu_cache_usage_perc * 100.0
        );

        assert_eq!(
            metrics.kv_active_blocks, 0,
            "Expected 0 active blocks, got {}",
            metrics.kv_active_blocks
836
837
838
        );
    }
}