scheduler.rs 35.9 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, KvStats, WorkerStats};
44
45
use crate::mocker::evictor::LRUEvictor;
use crate::mocker::kv_manager::KvManager;
46
use crate::mocker::protocols::{block_response_to_kv_event, MoveBlock, OutputSignal, PrefillCost};
47
use crate::mocker::protocols::{DirectRequest, MockEngineArgs, MoveBlockResponse};
48
use crate::mocker::sequence::ActiveSequence;
49
use crate::tokens::blocks::UniqueBlock;
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
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>,
67
68
    prefill: VecDeque<Uuid>,
    decode: LRUEvictor<Uuid>,
69
    requests: HashMap<Uuid, Request>,
70
71
72
73
    prefill_costs: HashMap<Uuid, PrefillCost>,
    max_num_batched_tokens: Option<usize>,
    active_tokens: usize,
    waiting_tokens: usize,
74
75
76
}

impl SchedulerState {
77
78
79
80
81
82
83
    fn new(max_num_batched_tokens: Option<usize>) -> Self {
        SchedulerState {
            max_num_batched_tokens,
            ..Default::default()
        }
    }

84
85
86
87
88
89
90
91
92
93
94
    /// 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)> {
95
96
97
98
99
        let uuid = self.waiting.pop_front()?;
        let request = self
            .requests
            .remove(&uuid)
            .expect("Request does not exist.");
100
101
102
        Some((uuid, request))
    }

103
104
105
106
107
108
    /// 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);
    }

109
    /// Move a UUID and its Request to the ready queue.
110
111
    fn move_to_prefill(&mut self, uuid: Uuid, active_seq: ActiveSequence, cost: PrefillCost) {
        self.waiting_tokens += cost.new_tokens;
112
        self.requests.insert(uuid, Request::Active(active_seq));
113
114
        self.prefill.push_back(uuid);
        self.prefill_costs.insert(uuid, cost);
115
116
    }

117
118
119
120
121
122
123
    /// 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
    /// - `is_full_prefill`: true if the entire sequence was prefilled, false if chunked
    fn try_prefill(&mut self) -> Option<(f64, Option<MoveBlock>, bool)> {
124
125
126
        let uuid = self.prefill.pop_front()?;

        // Remove and extract prefill_compute from prefill_costs
127
        let mut prefill_cost = self
128
129
130
            .prefill_costs
            .remove(&uuid)
            .expect("Expects valid prefill cost.");
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
165
166
167
168
169
170
171
        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 {
172
            panic!("Request does not exist.");
173
174
        };

175
176
177
178
179
180
181
182
183
184
        Some((
            prefill_compute,
            sequence.take_creation_signal(),
            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();
185
186
    }

187
188
189
190
191
192
193
194
    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)
195
196
    }

197
198
    fn num_active_requests(&self) -> usize {
        self.prefill.len() + self.decode.len()
199
200
201
202
    }

    /// Remove a UUID and its associated Request from collections.
    fn complete(&mut self, uuid: &Uuid) {
203
        tracing::debug!("Request {} will complete", uuid);
204
        self.decode.remove(uuid);
205
206
        self.requests.remove(uuid);
        self.prefill_costs.remove(uuid);
207
        self.active_tokens -= 1;
208
209
210
211
212
    }

    /// 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.
213
    fn preempt(&mut self) -> Vec<MoveBlock> {
214
        // Evict the oldest UUID from running
215
216
217
218
219
220
221
222
        let uuid = self
            .decode
            .evict()
            .expect("Nothing to evict for preemption.");
        let request = self
            .requests
            .remove(&uuid)
            .expect("Request does not exist.");
223
        self.prefill_costs.remove(&uuid);
224
225
        self.active_tokens -= 1;
        tracing::warn!("Request {uuid} will be preempted");
226

227
228
        // Reset the sequence and get the new sequence and signal
        // Insert the new sequence back into the requests map and add to waiting queue
229
230
231
232
233
        let Request::Active(mut active_sequence) = request else {
            panic!("Expected ActiveSequence in running queue")
        };
        let signals = active_sequence.reset_with_signal();

234
235
        // 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.
236

237
238
239
        self.first_in_line(uuid, Request::Active(active_sequence));

        signals
240
241
242
243
244
245
    }
}

/// Manages scheduling of requests using KvManager resources
#[derive(Clone)]
pub struct Scheduler {
246
    dp_rank: Option<u32>,
247
248
    state: Arc<Mutex<SchedulerState>>,
    kv_manager: Arc<Mutex<KvManager>>,
249
250
    request_tx: mpsc::UnboundedSender<DirectRequest>,
    hit_rates: Arc<Mutex<VecDeque<f32>>>,
251
252
253
254
255
}

impl Scheduler {
    /// Create a new Scheduler with the given parameters
    pub fn new(
256
257
258
259
        args: MockEngineArgs,
        dp_rank: Option<u32>,
        output_tx: Option<mpsc::UnboundedSender<OutputSignal>>,
        kv_events_tx: Option<mpsc::UnboundedSender<KvCacheEventData>>,
260
261
        cancellation_token: Option<CancellationToken>,
    ) -> Self {
262
        let state = Arc::new(Mutex::new(SchedulerState::new(args.max_num_batched_tokens)));
263

264
265
266
267
268
269
270
        // 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)
        };
271

272
273
274
275
276
277
278
279
280
281
282
283
284
        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
        );
285

286
287
        // Create channel for request handling
        let (request_tx, mut request_rx) = mpsc::unbounded_channel::<DirectRequest>();
288
289
290
291
292

        // 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();
293
294
        let cancel_token_clone = cancellation_token.unwrap_or_default().clone();
        let hit_rates_clone = hit_rates.clone();
295
296
297

        // Spawn main background task with cancellation token
        tokio::spawn(async move {
298
299
300
            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;
301
302
303
304
305
306
307
308
309
310
311

            loop {
                tokio::select! {
                    biased;

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

312
                    // Try Scheduling Requests - runs on normal interval or after simulation
313
                    _ = schedule_interval.tick() => {
314
315
316
317
318
                        // Skip if we just ran scheduling after simulation to prevent consecutive runs
                        if !should_schedule {
                            continue;
                        }

319
                        let mut state_guard = state_clone.lock().await;
320
                        let kv_manager_guard = kv_manager_clone.lock().await;
321
322
323

                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
                        // schedule anymore.
324
                        let mut current_blocks = kv_manager_guard.num_active_blocks();
325
                        let mut current_tokens = state_guard.active_tokens + state_guard.waiting_tokens;
326
327
                        let mut current_seqs = state_guard.num_active_requests();

328
                        while let Some((uuid, request)) = state_guard.next() {
329
                            let active_sequence = get_active_sequence(request, args.block_size, args.enable_prefix_caching);
330

331
332
333
                            // Update predictive budgets
                            let prefill_cost = kv_manager_guard.get_prefill_cost(&active_sequence);
                            let total_tokens = active_sequence.len();
334
335
                            // 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;
336
337
338
339
340
                            let new_tokens = prefill_cost.new_tokens;

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

342
                            // Check various budgets to see if possible to schedule
343
                            let under_block_budget = current_blocks as f64 <= (1. - args.watermark) * kv_manager_guard.max_capacity() as f64;
344
345
346
                            // 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);
347
348
349
350
351
                            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));
352
                                break;
353
354
355
356
357
358
359
360
361
362
363
                            }

                            // 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();
                                }
                            }
364

365
                            state_guard.move_to_prefill(uuid, active_sequence, prefill_cost);
366
                            should_schedule = false;
367
368
369
370
                        }
                    }

                    // Check for cancellation
371
                    _ = cancel_token_clone.cancelled() => {
372
373
374
375
376
377
378
379
                        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;

380
381
382
383
384
385
                        // 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
386
                        while let Some((prefill_compute, maybe_creation_signal, is_full_prefill)) = state_guard.try_prefill() {
387
388
389
                            // 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);
390

391
392
393
                            if let Some(creation_signal) = maybe_creation_signal {
                                if !process_signals(&mut kv_manager_guard, std::slice::from_ref(&creation_signal)) {
                                    panic!("Block allocation for prefilling cannot fail.");
394
                                }
395
396
397
398
399
400
401
402
403
404
405

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

                            // Impossible to schedule more prefills if we encounter one incomplete (chunked) prefill
                            if !is_full_prefill { break; }
406
                        }
407

408
409
                        state_guard.reset_active_tokens();

410
411
412
413
414
415
416
                        // 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;
                            };
417
418
419
420
421
422
                            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
423
                                for signal in state_guard.preempt() {
424
425
426
427
428
                                    kv_manager_guard.process(&signal);
                                }
                                continue;
                            }

429
430
431
432
433
                            // 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));
                                }
434
435
                            }

436
437
438
439
440
441
442
443
444
                            // 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()
                                });
445
446
                            }

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

                            if send_failed || is_complete {
                                state_guard.complete(&uuid);
                                continue;
456
457
458
                            }
                        }

459
460
461
462
463
464
                        // 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;
465
466
467
468
469
470
471
                        }
                    }
                }
            }
        });

        Self {
472
            dp_rank,
473
474
475
            state,
            kv_manager,
            request_tx,
476
            hit_rates,
477
478
479
480
481
        }
    }

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

    pub fn request_sender(&self) -> mpsc::UnboundedSender<DirectRequest> {
        self.request_tx.clone()
487
488
489
490
491
492
493
494
495
    }

    pub async fn waiting_count(&self) -> usize {
        let state = self.state.lock().await;
        state.waiting.len()
    }

    pub async fn running_count(&self) -> usize {
        let state = self.state.lock().await;
496
        state.decode.len()
497
498
    }

499
500
501
502
503
504
505
506
507
508
    pub async fn waiting_tokens(&self) -> usize {
        let state = self.state.lock().await;
        state.waiting_tokens
    }

    pub async fn active_tokens(&self) -> usize {
        let state = self.state.lock().await;
        state.active_tokens
    }

509
510
511
512
513
514
515
    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 {
516
        // Acquire all locks in consistent order: state -> kv_manager -> hit_rates
517
518
        let state = self.state.lock().await;
        let kv_manager = self.kv_manager.lock().await;
519
520
521
522
523
        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;
524

525
        // Get KV manager metrics
526
527
528
529
530
531
532
533
        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
        };

534
535
536
537
538
539
540
541
        // 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
        };

542
        let worker_stats = WorkerStats {
543
544
            data_parallel_rank: self.dp_rank,
            request_active_slots,
545
546
547
548
549
            request_total_slots: 1024, // vllm max_num_seqs for gpu >= 70 vram, otherwise 256, fallback is 128
            num_requests_waiting,
        };

        let kv_stats = KvStats {
550
551
552
            kv_active_blocks: active_blocks_count,
            kv_total_blocks: total_capacity,
            gpu_cache_usage_perc,
553
            gpu_prefix_cache_hit_rate,
554
555
556
557
558
559
560
561
        };

        let spec_decode_stats = None;

        ForwardPassMetrics {
            worker_stats,
            kv_stats,
            spec_decode_stats,
562
        }
563
        // Guards drop naturally here in reverse order (LIFO): hit_rates_guard, kv_manager, state
564
565
566
567
    }
}

/// Convert a Request to an ActiveSequence
568
569
570
571
572
fn get_active_sequence(
    request: Request,
    block_size: usize,
    enable_prefix_caching: bool,
) -> ActiveSequence {
573
574
575
576
577
578
579
580
581
582
583
584
    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),
585
        enable_prefix_caching,
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
    )
}

/// 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
606
        let MoveBlock::Use(blocks) = signal else {
607
            panic!("Failed signal is Invalid. Has to fail on generation signal, but failed on {signal:?}");
608
609
610
        };

        // Verify the signal contains exactly one block
611
612
613
614
615
616
        let num_blocks = blocks.len();
        let num_active_blocks = kv_manager_guard.num_active_blocks();
        if num_blocks != 1 {
            panic!(
                "Failed signal is Invalid. Tried to create (prefill) {num_blocks} blocks on top of {num_active_blocks} active blocks."
            );
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
        }

        // 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]
637
638
639
640
641
642
643
644
    #[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)]
645
    #[tokio::test]
646
647
648
    async fn test_scheduler_token_generation_patterns(
        #[case] use_shared_tokens: bool,
        #[case] enable_prefix_caching: bool,
649
        #[case] enable_chunked_prefill: bool,
650
    ) {
651
652
653
        std::env::set_var("RUST_LOG", "debug");

        let kv_capacity: usize = 500;
654
        let block_size: usize = 64;
655
        let num_requests: usize = 200;
656
657
658
659
        let input_len: usize = 1000;
        let max_output_tokens: usize = 100;

        // Create channel for token output
660
661
662
663
664
665
666
667
        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)
668
            .enable_chunked_prefill(enable_chunked_prefill)
669
670
671
672
673
            .build()
            .unwrap();

        // Create scheduler with new args struct
        let scheduler = Scheduler::new(args, None, Some(output_tx), None, None);
674
675
676
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

        // 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,
704
                dp_rank: None,
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
            };
            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;
729
                    println!("Forward Pass Metrics: {_metrics:#?}");
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
                }

                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!(
748
            "Test completed in: {elapsed:?} for {} case with prefix_caching={enable_prefix_caching} and chunked_prefill={enable_chunked_prefill}",
749
750
751
752
            if use_shared_tokens {
                "caching"
            } else {
                "random"
753
            }
754
755
756
757
        );

        // Assert that we received the expected number of tokens
        assert!(
758
759
760
            received_tokens == expected_tokens,
            "Received {received_tokens} tokens but expected exactly {expected_tokens}"
        );
761
762
763
764
765
766
767
768
769
770
771
772

        let active_tokens = scheduler.active_tokens().await;
        assert!(
            active_tokens == 0,
            "Scheduler still have {active_tokens} active tokens but expected 0"
        );

        let waiting_tokens = scheduler.waiting_tokens().await;
        assert!(
            waiting_tokens == 0,
            "Scheduler still have {waiting_tokens} waiting tokens but expected 0"
        );
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
836
837
838
839
840
841
842
843
844
845
846
847
848
849
    }

    #[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!(
850
            metrics.worker_stats.num_requests_waiting, 0,
851
            "Expected no waiting requests, got {}",
852
            metrics.worker_stats.num_requests_waiting
853
854
855
        );

        assert!(
856
            metrics.kv_stats.gpu_prefix_cache_hit_rate > 0.8,
857
            "Expected cache hit rate > 0.8, got {}",
858
            metrics.kv_stats.gpu_prefix_cache_hit_rate
859
860
861
862
        );

        println!(
            "Test passed! Cache hit rate: {:.3}",
863
            metrics.kv_stats.gpu_prefix_cache_hit_rate
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
        );
        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!(
919
            metrics.kv_stats.gpu_cache_usage_perc,
920
921
            0.0,
            "Expected GPU cache usage to be 0%, got {}%",
922
            metrics.kv_stats.gpu_cache_usage_perc * 100.0
923
924
925
        );

        assert_eq!(
926
            metrics.kv_stats.kv_active_blocks, 0,
927
            "Expected 0 active blocks, got {}",
928
            metrics.kv_stats.kv_active_blocks
929
930
931
        );
    }
}