vllm.rs 32 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
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-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
32
33
use crate::common::evictor::LRUEvictor;
use crate::common::perf_model::PerfModel;
use crate::common::protocols::{
34
35
    DirectRequest, KvCacheEventSink, MockEngineArgs, MoveBlock, OutputSignal, PrefillCost,
    WorkerType,
Yan Ru Pei's avatar
Yan Ru Pei committed
36
};
37
38
use crate::common::running_mean::RunningMean;
use crate::common::sequence::ActiveSequence;
39
use crate::common::utils::sleep_until_precise;
40
use crate::kv_manager::KvManager;
41
use dynamo_kv_router::protocols::DpRank;
42
use dynamo_tokens::blocks::UniqueBlock;
Yan Ru Pei's avatar
Yan Ru Pei committed
43
use std::collections::{HashMap, VecDeque};
44
use std::sync::Arc;
45
use std::time::Instant;
46
use tokio::sync::mpsc;
47
use tokio::time::Duration;
48
49
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
50
use validator::Validate;
51

52
53
54
55
56
57
58
/// Simple metrics struct for mocker's internal use
#[derive(Clone, Default, Debug)]
pub struct MockerMetrics {
    pub dp_rank: DpRank,
    pub active_decode_blocks: u64,
}

59
60
61
62
63
64
65
66
67
/// 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
    requests: HashMap<Uuid, Request>,
71
72
73
74
    prefill_costs: HashMap<Uuid, PrefillCost>,
    max_num_batched_tokens: Option<usize>,
    active_tokens: usize,
    waiting_tokens: usize,
75
76
77
}

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

85
86
87
88
    fn is_empty(&self) -> bool {
        self.requests.is_empty()
    }

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

108
109
110
111
112
113
    /// 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);
    }

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

122
123
124
125
126
127
    /// 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
128
    fn try_prefill(&mut self, perf_model: &PerfModel) -> Option<(f64, Option<MoveBlock>, bool)> {
129
130
131
        let uuid = self.prefill.pop_front()?;

        // Remove and extract prefill_compute from prefill_costs
132
        let mut prefill_cost = self
133
134
135
            .prefill_costs
            .remove(&uuid)
            .expect("Expects valid prefill cost.");
136

137
138
139
140
141
142
143
144
145
146
147
148
149
        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
        {
150
151
            let prefill_compute =
                prefill_cost.predict_prefill_compute(Some(prefill_tokens), perf_model);
152
153
            prefill_cost.new_tokens -= prefill_tokens;
            assert!(
154
155
                prefill_cost.new_tokens > 0,
                "Encountered negative prefill tokens."
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
            );

            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;

172
            (prefill_cost.predict_prefill_compute(None, perf_model), true)
173
174
175
176
177
        };

        // 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 {
178
            panic!("Request does not exist.");
179
180
        };

181
182
183
184
185
186
187
188
189
190
        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();
191
192
    }

193
194
195
196
197
198
199
200
    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)
201
202
    }

203
204
    fn num_active_requests(&self) -> usize {
        self.prefill.len() + self.decode.len()
205
206
207
208
    }

    /// Remove a UUID and its associated Request from collections.
    fn complete(&mut self, uuid: &Uuid) {
209
        tracing::trace!("Request {uuid} will complete");
210
        self.decode.remove(uuid);
211
212
        self.requests.remove(uuid);
        self.prefill_costs.remove(uuid);
213
        self.active_tokens -= 1;
214
215
216
217
218
    }

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

233
234
        // Reset the sequence and get the new sequence and signal
        // Insert the new sequence back into the requests map and add to waiting queue
235
236
237
238
239
        let Request::Active(mut active_sequence) = request else {
            panic!("Expected ActiveSequence in running queue")
        };
        let signals = active_sequence.reset_with_signal();

240
241
        // 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.
242

243
244
245
        self.first_in_line(uuid, Request::Active(active_sequence));

        signals
246
247
248
    }
}

249
250
251
252
253
254
255
256
257
258
/// Cancels its token when dropped. Shared via Arc so the background task is
/// only cancelled when the last Scheduler clone is dropped.
struct CancelGuard(CancellationToken);

impl Drop for CancelGuard {
    fn drop(&mut self) {
        self.0.cancel();
    }
}

259
260
261
/// Manages scheduling of requests using KvManager resources
#[derive(Clone)]
pub struct Scheduler {
262
    request_tx: mpsc::UnboundedSender<DirectRequest>,
263
    metrics_rx: tokio::sync::watch::Receiver<MockerMetrics>,
264
    _cancel_guard: Arc<CancelGuard>,
265
266
267
268
269
}

impl Scheduler {
    /// Create a new Scheduler with the given parameters
    pub fn new(
270
        args: MockEngineArgs,
Yan Ru Pei's avatar
Yan Ru Pei committed
271
        dp_rank: u32,
272
        output_tx: Option<mpsc::UnboundedSender<OutputSignal>>,
273
        kv_event_sink: Option<Arc<dyn KvCacheEventSink>>,
274
275
        cancellation_token: Option<CancellationToken>,
    ) -> Self {
276
        args.validate().expect("invalid MockEngineArgs");
277

278
279
        // Create channel for request handling
        let (request_tx, mut request_rx) = mpsc::unbounded_channel::<DirectRequest>();
280
281
282
283
        let initial_metrics = MockerMetrics {
            dp_rank,
            active_decode_blocks: 0,
        };
284
        let (metrics_tx, metrics_rx) =
285
            tokio::sync::watch::channel::<MockerMetrics>(initial_metrics);
286

287
288
289
        let cancel_token = cancellation_token.unwrap_or_default();
        let cancel_token_clone = cancel_token.clone();
        let cancel_guard = Arc::new(CancelGuard(cancel_token));
290
291
292

        // Spawn main background task with cancellation token
        tokio::spawn(async move {
293
294
            // Create state and kv_manager as local variables owned by this task
            let mut state = SchedulerState::new(args.max_num_batched_tokens);
295
            let mut kv_manager = KvManager::new_with_event_sink(
Yan Ru Pei's avatar
Yan Ru Pei committed
296
297
                args.num_gpu_blocks,
                args.block_size,
298
                kv_event_sink,
Yan Ru Pei's avatar
Yan Ru Pei committed
299
300
301
                dp_rank,
            );
            let mut hit_rates = RunningMean::new(1000);
302
303

            loop {
Yan Ru Pei's avatar
Yan Ru Pei committed
304
                // 1. Receive requests
305
306
307
308
309
                if receive_requests(&mut state, &mut request_rx, &cancel_token_clone)
                    .await
                    .is_none()
                {
                    break;
310
                }
311

Yan Ru Pei's avatar
Yan Ru Pei committed
312
313
314
315
                // 2. Schedule waiting requests (once per iteration)
                try_schedule(&mut state, &kv_manager, &mut hit_rates, &args);

                // 3. Simulate prefill + decode
316
                simulate_prefill(
317
318
319
320
                    &mut state,
                    &mut kv_manager,
                    &args.perf_model,
                    args.worker_type,
321
322
323
                    args.speedup_ratio,
                )
                .await;
324

325
                simulate_decode(
326
327
328
329
330
                    &mut state,
                    &mut kv_manager,
                    &output_tx,
                    &args.perf_model,
                    args.block_size,
331
332
333
                    args.speedup_ratio,
                )
                .await;
334
335

                // 4. Send metrics once per forward pass (after all prefill and decode processing)
336
                let _ = metrics_tx.send(MockerMetrics {
337
                    dp_rank,
338
339
                    active_decode_blocks: kv_manager.num_active_blocks() as u64,
                });
340
341
342
343
344
            }
        });

        Self {
            request_tx,
345
            metrics_rx,
346
            _cancel_guard: cancel_guard,
347
348
349
350
351
        }
    }

    /// Add a new request to the waiting queue
    pub async fn receive(&self, request: DirectRequest) {
352
353
354
355
356
        let _ = self.request_tx.send(request);
    }

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

359
    /// Get a watch receiver for forward pass metrics
360
    pub fn metrics_receiver(&self) -> tokio::sync::watch::Receiver<MockerMetrics> {
361
362
363
        self.metrics_rx.clone()
    }
}
364

365
366
367
368
369
370
371
372
373
374
375
376
/// Receive requests from the channel.
/// Returns `Some(())` to continue the loop, `None` to break (on cancellation).
async fn receive_requests(
    state: &mut SchedulerState,
    request_rx: &mut mpsc::UnboundedReceiver<DirectRequest>,
    cancel_token: &CancellationToken,
) -> Option<()> {
    if cancel_token.is_cancelled() {
        return None;
    }

    if state.is_empty() {
377
        // Fully idle - block until new request arrives or shutdown
378
379
380
381
382
        tokio::select! {
            biased;
            _ = cancel_token.cancelled() => {
                return None;
            }
383
384
385
386
            result = request_rx.recv() => {
                let Some(request) = result else {
                    return None; // channel closed
                };
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
                state.receive(request);
                return Some(());
            }
        }
    }

    // Has active/waiting work - collect any pending requests without blocking
    while let Ok(request) = request_rx.try_recv() {
        state.receive(request);
    }

    Some(())
}

/// Simulate prefill phase for all pending prefill requests.
/// Returns the total prefill compute time.
403
async fn simulate_prefill(
404
405
406
407
    state: &mut SchedulerState,
    kv_manager: &mut KvManager,
    perf_model: &PerfModel,
    worker_type: WorkerType,
408
    speedup_ratio: f64,
409
) -> Duration {
410
    let start_time = Instant::now();
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
    let mut total_time = Duration::ZERO;

    while let Some((prefill_compute, maybe_creation_signal, is_full_prefill)) =
        state.try_prefill(perf_model)
    {
        // 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.
        // For decode workers, skip adding prefill compute time
        if worker_type != WorkerType::Decode {
            total_time += Duration::from_secs_f64(prefill_compute / 1000.0);
        }

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

        // Impossible to schedule more prefills if we encounter one incomplete (chunked) prefill
        if !is_full_prefill {
            break;
        }
    }
434
435
436
437
438

    if speedup_ratio > 0.0 && total_time > Duration::ZERO {
        let sleep_duration = Duration::from_secs_f64(total_time.as_secs_f64() / speedup_ratio);
        let deadline = start_time + sleep_duration;

439
        sleep_until_precise(deadline).await;
440
    }
441

442
443
444
445
446
    total_time
}

/// Simulate decode phase for all active decode requests.
/// Returns the total decode compute time.
447
async fn simulate_decode(
448
449
450
451
452
    state: &mut SchedulerState,
    kv_manager: &mut KvManager,
    output_tx: &Option<mpsc::UnboundedSender<OutputSignal>>,
    perf_model: &PerfModel,
    block_size: usize,
453
    speedup_ratio: f64,
454
) -> Duration {
455
456
    let start_time = Instant::now();

457
458
    // Compute decode timing
    let active_kv_tokens = kv_manager.num_active_blocks() * block_size;
459

460
    // Compute average context length across all active decode requests
461
    let total_length: usize = state
462
463
        .decode
        .keys()
464
465
466
        .map(|uuid| {
            if let Request::Active(seq) = state.requests.get(uuid).unwrap() {
                seq.len()
467
            } else {
468
                0
469
            }
470
471
472
473
        })
        .sum();
    let count = state.decode.len();

474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
    let context_length = if count > 0 { total_length / count } else { 0 };
    let decoding_time = perf_model.predict_decode_time(active_kv_tokens, context_length);
    let total_time = Duration::from_secs_f64(decoding_time / 1000.0);

    state.reset_active_tokens();

    // Process decoding
    let uuids: Vec<Uuid> = state.decode.keys().cloned().collect();
    for uuid in uuids {
        let Some(sequence) = state.run(uuid) else {
            continue;
        };
        let signals = sequence.generate();

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

        // Check completion and send notification
        let is_complete = sequence.generated_tokens() >= sequence.max_output_tokens();

501
502
503
504
505
506
507
        let send_failed = output_tx.as_ref().is_some_and(|tx| {
            tx.send(OutputSignal {
                uuid,
                completed: is_complete,
            })
            .is_err()
        });
508
509
510
511
512
513
514
515
516
517
518

        if send_failed {
            for signal in &sequence.free_signal() {
                kv_manager.process(signal);
            }
        }

        if send_failed || is_complete {
            state.complete(&uuid);
        }
    }
519
520
521
522
523

    if speedup_ratio > 0.0 && total_time > Duration::ZERO {
        let sleep_duration = Duration::from_secs_f64(total_time.as_secs_f64() / speedup_ratio);
        let deadline = start_time + sleep_duration;

524
        sleep_until_precise(deadline).await;
525
    }
526

527
528
529
    total_time
}

Yan Ru Pei's avatar
Yan Ru Pei committed
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
/// Attempts to schedule waiting requests from the state queue.
/// Returns the number of requests successfully scheduled.
fn try_schedule(
    state: &mut SchedulerState,
    kv_manager: &KvManager,
    hit_rates: &mut RunningMean<f32>,
    args: &MockEngineArgs,
) -> usize {
    let mut scheduled_count = 0;
    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();

    while let Some((uuid, request)) = state.next() {
        // Convert Request to ActiveSequence
        let active_sequence = match request {
            Request::Active(active_seq) => active_seq,
            Request::Direct(direct_request) => ActiveSequence::new(
                direct_request.tokens,
                direct_request.max_output_tokens,
                Some(args.block_size),
                args.enable_prefix_caching,
552
                args.zmq_kv_events_port.is_some(),
Yan Ru Pei's avatar
Yan Ru Pei committed
553
554
            ),
        };
555

Yan Ru Pei's avatar
Yan Ru Pei committed
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
        // Update predictive budgets
        let prefill_cost = kv_manager.get_prefill_cost(&active_sequence);
        let total_tokens = active_sequence.len();
        // 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;
        let new_tokens = prefill_cost.new_tokens;

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

        // Check various budgets to see if possible to schedule
        let under_block_budget =
            current_blocks as f64 <= (1. - args.watermark) * kv_manager.max_capacity() as f64;
        // 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);
        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.first_in_line(uuid, Request::Active(active_sequence));
            break;
        }

        // 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
        };
        hit_rates.push(hit_rate);

        state.move_to_prefill(uuid, active_sequence, prefill_cost);
        scheduled_count += 1;
    }
598

Yan Ru Pei's avatar
Yan Ru Pei committed
599
    scheduled_count
600
601
602
603
604
605
606
607
608
}

/// 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.
609
fn process_signals(kv_manager: &mut KvManager, signals: &[MoveBlock]) -> bool {
610
    for signal in signals {
611
        if kv_manager.process(signal) {
612
613
614
615
            continue;
        }

        // Check we have a Use signal with blocks
616
        let MoveBlock::Use(blocks, _hashes, ..) = signal else {
617
618
619
            panic!(
                "Failed signal is Invalid. Has to fail on generation signal, but failed on {signal:?}"
            );
620
621
622
        };

        // Verify the signal contains exactly one block
623
        let num_blocks = blocks.len();
624
        let num_active_blocks = kv_manager.num_active_blocks();
625
626
627
628
        if num_blocks != 1 {
            panic!(
                "Failed signal is Invalid. Tried to create (prefill) {num_blocks} blocks on top of {num_active_blocks} active blocks."
            );
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
        }

        // 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;
647
    use tokio::time::interval;
648

649
650
    /// Helper function to verify that the scheduler is idle (no active KV blocks)
    fn assert_scheduler_idle(metrics: &MockerMetrics) {
651
        assert_eq!(
652
            metrics.active_decode_blocks, 0,
653
            "Expected 0 active blocks, got {}",
654
            metrics.active_decode_blocks
655
656
657
        );
    }

658
    #[rstest]
659
660
661
662
663
664
665
666
    #[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)]
667
    #[tokio::test]
668
669
670
    async fn test_scheduler_token_generation_patterns(
        #[case] use_shared_tokens: bool,
        #[case] enable_prefix_caching: bool,
671
        #[case] enable_chunked_prefill: bool,
672
    ) {
673
        unsafe { std::env::set_var("RUST_LOG", "debug") };
674
675

        let kv_capacity: usize = 500;
676
        let block_size: usize = 64;
677
        let num_requests: usize = 200;
678
679
680
681
        let input_len: usize = 1000;
        let max_output_tokens: usize = 100;

        // Create channel for token output
682
683
684
685
686
687
688
689
        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)
690
            .enable_chunked_prefill(enable_chunked_prefill)
691
692
693
694
            .build()
            .unwrap();

        // Create scheduler with new args struct
Yan Ru Pei's avatar
Yan Ru Pei committed
695
        let scheduler = Scheduler::new(args, 0, Some(output_tx), None, None);
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

        // 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
726
                dp_rank: 0,
727
728
729
730
731
732
733
734
735
736
737
738
739
740
            };
            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);

741
742
743
        // Get metrics receiver
        let metrics_rx = scheduler.metrics_receiver();

744
745
746
747
748
749
750
751
752
        // 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() => {
753
                    let _metrics = metrics_rx.borrow().clone();
754
                    tracing::debug!("Forward Pass Metrics: {_metrics:#?}");
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
                }

                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!(
773
            "Test completed in: {elapsed:?} for {} case with prefix_caching={enable_prefix_caching} and chunked_prefill={enable_chunked_prefill}",
774
775
776
777
            if use_shared_tokens {
                "caching"
            } else {
                "random"
778
            }
779
780
781
782
        );

        // Assert that we received the expected number of tokens
        assert!(
783
784
785
            received_tokens == expected_tokens,
            "Received {received_tokens} tokens but expected exactly {expected_tokens}"
        );
786

787
788
        // Wait a bit for final metrics update to propagate
        tokio::time::sleep(Duration::from_millis(100)).await;
789

790
791
        let metrics = scheduler.metrics_receiver().borrow().clone();
        assert_scheduler_idle(&metrics);
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
    }

    #[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
814
        let scheduler = Scheduler::new(args, 0, Some(output_tx), None, None);
815
816
817
818
819
820
821
822
823
824

        // 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
825
                dp_rank: 0,
826
827
828
829
830
831
832
833
834
835
836
837
838
            };
            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);

839
840
841
        // Get metrics receiver
        let metrics_rx = scheduler.metrics_receiver();

842
843
844
845
846
847
848
849
850
        // 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() => {
851
                    let _metrics = metrics_rx.borrow().clone();
852
                    tracing::debug!("Forward Pass Metrics: {_metrics:#?}");
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
                }

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

868
869
870
        // Wait a bit for final metrics update
        tokio::time::sleep(Duration::from_millis(100)).await;

871
        // Verify forward pass metrics - scheduler should be idle after completing all requests
872
        let metrics = metrics_rx.borrow().clone();
873
        assert_scheduler_idle(&metrics);
874

875
        println!("Test passed! Received {received_tokens} tokens");
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
    }

    #[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
896
        let scheduler = Scheduler::new(args, 0, Some(output_tx), None, None);
897
898
899
900
901
902
903

        // 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
904
            dp_rank: 0,
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
        };

        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
926
927
        let metrics_rx = scheduler.metrics_receiver();
        let metrics = metrics_rx.borrow().clone();
928

929
        assert_scheduler_idle(&metrics);
930
931
    }
}