## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/realizar/src/scheduler/mod.rs

    Line
    Count
    Source
    1
    //! Continuous Batching Scheduler
    2
    //!
    3
    //! Per spec §8: Implements continuous batching for LLM serving based on vLLM.
    4
    //! Reference: [8] Yu et al. (2022) "Orca: A Distributed Serving System for Transformer-Based Generative Models"
    5
    //!
    6
    //! ## Key Features
    7
    //!
    8
    //! - **Iteration-Level Scheduling**: New requests join batch at any iteration
    9
    //! - **Preemption**: Low-priority requests can be preempted for high-priority
    10
    //! - **Memory-Aware**: Respects KV cache limits when scheduling
    11
    //! - **Fair Queuing**: Prevents starvation of long requests
    12
    //!
    13
    //! ## Scheduling Algorithm
    14
    //!
    15
    //! ```text
    16
    //! while running:
    17
    //!   1. Check for completed sequences (EOS or max_tokens)
    18
    //!   2. Preempt sequences if memory pressure
    19
    //!   3. Schedule waiting sequences if space available
    20
    //!   4. Run one iteration of generation for batch
    21
    //! ```
    22
    23
    // Module-level clippy allows
    24
    #![allow(clippy::must_use_candidate)]
    25
    #![allow(clippy::return_self_not_must_use)]
    26
    #![allow(clippy::missing_errors_doc)]
    27
    #![allow(clippy::unnecessary_wraps)] // Result wrapping for API consistency
    28
    #![allow(clippy::derivable_impls)] // Manual impl for documentation clarity
    29
    #![allow(clippy::option_if_let_else)] // map_or is more readable
    30
    31
    use crate::paged_kv::{PagedCacheError, PagedKvCache, SeqId};
    32
    use serde::{Deserialize, Serialize};
    33
    use std::collections::{BinaryHeap, HashMap, VecDeque};
    34
    use std::time::Instant;
    35
    use thiserror::Error;
    36
    37
    // PMAT-802: Extracted modules
    38
    mod chunked_prefill;
    39
    mod types;
    40
    pub use chunked_prefill::{ChunkedPrefillConfig, ChunkedPrefillScheduler, ChunkedPrefillState, ChunkedPrefillStats};
    41
    pub use types::{Priority, SequenceState, SchedulerStats};
    42
    43
    /// Error type for scheduler operations
    44
    #[derive(Debug, Error)]
    45
    pub enum SchedulerError {
    46
        /// Queue is full
    47
        #[error("Request queue full: capacity {capacity}")]
    48
        QueueFull {
    49
            /// Queue capacity
    50
            capacity: usize,
    51
        },
    52
    53
        /// Request not found
    54
        #[error("Request not found: {0}")]
    55
        RequestNotFound(u64),
    56
    57
        /// KV cache error
    58
        #[error("KV cache error: {0}")]
    59
        CacheError(#[from] PagedCacheError),
    60
    61
        /// Invalid state
    62
        #[error("Invalid scheduler state: {0}")]
    63
        InvalidState(String),
    64
    }
    65
    66
    /// Generation request
    67
    #[derive(Debug, Clone)]
    68
    pub struct SchedulerRequest {
    69
        /// Unique request ID
    70
        pub request_id: u64,
    71
        /// Input token IDs
    72
        pub input_ids: Vec<u32>,
    73
        /// Maximum tokens to generate
    74
        pub max_tokens: usize,
    75
        /// Priority level
    76
        pub priority: Priority,
    77
        /// Arrival time
    78
        pub arrival_time: Instant,
    79
        /// Sequence ID (assigned when scheduled)
    80
        pub seq_id: Option<SeqId>,
    81
        /// Current state
    82
        pub state: SequenceState,
    83
        /// Generated tokens so far
    84
        pub generated_tokens: Vec<u32>,
    85
        /// Number of decode iterations
    86
        pub iterations: usize,
    87
    }
    88
    89
    impl SchedulerRequest {
    90
        /// Create a new request
    91
    22
        pub fn new(request_id: u64, input_ids: Vec<u32>, max_tokens: usize) -> Self {
    92
    22
            Self {
    93
    22
                request_id,
    94
    22
                input_ids,
    95
    22
                max_tokens,
    96
    22
                priority: Priority::default(),
    97
    22
                arrival_time: Instant::now(),
    98
    22
                seq_id: None,
    99
    22
                state: SequenceState::Waiting,
    100
    22
                generated_tokens: Vec::new(),
    101
    22
                iterations: 0,
    102
    22
            }
    103
    22
        }
    104
    105
        /// Set priority
    106
    16
        pub fn with_priority(mut self, priority: Priority) -> Self {
    107
    16
            self.priority = priority;
    108
    16
            self
    109
    16
        }
    110
    111
        /// Total tokens (input + generated)
    112
    3
        pub fn total_tokens(&self) -> usize {
    113
    3
            self.input_ids.len() + self.generated_tokens.len()
    114
    3
        }
    115
    116
        /// Remaining tokens to generate
    117
    2
        pub fn remaining_tokens(&self) -> usize {
    118
    2
            self.max_tokens.saturating_sub(self.generated_tokens.len())
    119
    2
        }
    120
    121
        /// Check if generation is complete
    122
    6
        pub fn is_complete(&self, eos_token: u32) -> bool {
    123
    6
            self.generated_tokens.len() >= self.max_tokens
    124
    5
                || self.generated_tokens.last() == Some(&eos_token)
    125
    6
        }
    126
    127
        /// Time waiting in queue
    128
    3
        pub fn wait_time(&self) -> std::time::Duration {
    129
    3
            self.arrival_time.elapsed()
    130
    3
        }
    131
    }
    132
    133
    /// Priority-aware entry for the waiting queue
    134
    #[derive(Debug)]
    135
    struct PriorityEntry {
    136
        priority: Priority,
    137
        arrival_time: Instant,
    138
        request_id: u64,
    139
    }
    140
    141
    impl PartialEq for PriorityEntry {
    142
    2
        fn eq(&self, other: &Self) -> bool {
    143
    2
            self.request_id == other.request_id
    144
    2
        }
    145
    }
    146
    147
    impl Eq for PriorityEntry {}
    148
    149
    impl PartialOrd for PriorityEntry {
    150
    6
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
    151
    6
            Some(self.cmp(other))
    152
    6
        }
    153
    }
    154
    155
    impl Ord for PriorityEntry {
    156
    6
        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
    157
            // Higher priority first, then earlier arrival
    158
    6
            match self.priority.cmp(&other.priority) {
    159
    4
                std::cmp::Ordering::Equal => other.arrival_time.cmp(&self.arrival_time),
    160
    2
                other => other,
    161
            }
    162
    6
        }
    163
    }
    164
    165
    /// Scheduler output for one iteration
    166
    #[derive(Debug, Clone, Default)]
    167
    pub struct SchedulerOutput {
    168
        /// Sequences to run this iteration
    169
        pub scheduled_seq_ids: Vec<SeqId>,
    170
        /// Request IDs for scheduled sequences
    171
        pub scheduled_request_ids: Vec<u64>,
    172
        /// Sequences that were preempted
    173
        pub preempted_seq_ids: Vec<SeqId>,
    174
        /// Sequences that completed
    175
        pub completed_request_ids: Vec<u64>,
    176
        /// Number of tokens in prefill phase
    177
        pub num_prefill_tokens: usize,
    178
        /// Number of tokens in decode phase
    179
        pub num_decode_tokens: usize,
    180
    }
    181
    182
    impl SchedulerOutput {
    183
        /// Total tokens scheduled
    184
    2
        pub fn total_tokens(&self) -> usize {
    185
    2
            self.num_prefill_tokens + self.num_decode_tokens
    186
    2
        }
    187
    188
        /// Check if batch is empty
    189
    2
        pub fn is_empty(&self) -> bool {
    190
    2
            self.scheduled_seq_ids.is_empty()
    191
    2
        }
    192
    }
    193
    194
    /// Continuous batching scheduler
    195
    pub struct Scheduler {
    196
        /// All requests by ID
    197
        requests: HashMap<u64, SchedulerRequest>,
    198
        /// Waiting queue (priority-ordered)
    199
        waiting_queue: BinaryHeap<PriorityEntry>,
    200
        /// Running sequences
    201
        running: Vec<u64>,
    202
        /// Preempted sequences (can be resumed)
    203
        preempted: VecDeque<u64>,
    204
        /// Maximum batch size
    205
        max_batch_size: usize,
    206
        /// Maximum queue size
    207
        max_queue_size: usize,
    208
        /// Maximum tokens per batch
    209
        max_tokens_per_batch: usize,
    210
        /// Next request ID
    211
        next_request_id: u64,
    212
        /// Statistics
    213
        stats: SchedulerStats,
    214
        /// Total wait time for completed requests (for averaging)
    215
        total_wait_time_ms: f64,
    216
    }
    217
    218
    impl Scheduler {
    219
        /// Create a new scheduler
    220
    16
        pub fn new(max_batch_size: usize, max_queue_size: usize) -> Self {
    221
    16
            Self {
    222
    16
                requests: HashMap::new(),
    223
    16
                waiting_queue: BinaryHeap::new(),
    224
    16
                running: Vec::new(),
    225
    16
                preempted: VecDeque::new(),
    226
    16
                max_batch_size,
    227
    16
                max_queue_size,
    228
    16
                max_tokens_per_batch: max_batch_size * 2048, // Default: assume 2k context
    229
    16
                next_request_id: 0,
    230
    16
                stats: SchedulerStats::default(),
    231
    16
                total_wait_time_ms: 0.0,
    232
    16
            }
    233
    16
        }
    234
    235
        /// Set maximum tokens per batch
    236
    1
        pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
    237
    1
            self.max_tokens_per_batch = max_tokens;
    238
    1
            self
    239
    1
        }
    240
    241
        /// Add a new request to the queue
    242
    12
        pub fn add_request(
    243
    12
            &mut self,
    244
    12
            input_ids: Vec<u32>,
    245
    12
            max_tokens: usize,
    246
    12
        ) -> Result<u64, SchedulerError> {
    247
    12
            self.add_request_with_priority(input_ids, max_tokens, Priority::Normal)
    248
    12
        }
    249
    250
        /// Add a request with priority
    251
    16
        pub fn add_request_with_priority(
    252
    16
            &mut self,
    253
    16
            input_ids: Vec<u32>,
    254
    16
            max_tokens: usize,
    255
    16
            priority: Priority,
    256
    16
        ) -> Result<u64, SchedulerError> {
    257
    16
            if self.waiting_queue.len() >= self.max_queue_size {
    258
    1
                return Err(SchedulerError::QueueFull {
    259
    1
                    capacity: self.max_queue_size,
    260
    1
                });
    261
    15
            }
    262
    263
    15
            let request_id = self.next_request_id;
    264
    15
            self.next_request_id += 1;
    265
    266
    15
            let request =
    267
    15
                SchedulerRequest::new(request_id, input_ids, max_tokens).with_priority(priority);
    268
    15
            let entry = PriorityEntry {
    269
    15
                priority,
    270
    15
                arrival_time: request.arrival_time,
    271
    15
                request_id,
    272
    15
            };
    273
    274
    15
            self.requests.insert(request_id, request);
    275
    15
            self.waiting_queue.push(entry);
    276
    15
            self.stats.total_requests += 1;
    277
    15
            self.stats.queue_depth = self.waiting_queue.len();
    278
    279
    15
            Ok(request_id)
    280
    16
        }
    281
    282
        /// Schedule one iteration of generation
    283
    12
        pub fn schedule(
    284
    12
            &mut self,
    285
    12
            kv_cache: &mut PagedKvCache,
    286
    12
            eos_token: u32,
    287
    12
        ) -> Result<SchedulerOutput, SchedulerError> {
    288
    12
            let mut output = SchedulerOutput::default();
    289
    290
            // 1. Check for completed sequences
    291
    12
            self.check_completions(&mut output, eos_token);
    292
    293
            // 2. Preempt if memory pressure (simplified: check if we can fit new sequences)
    294
    12
            self.handle_preemption(&mut output, kv_cache);
    295
    296
            // 3. Resume preempted sequences if possible
    297
    12
            self.resume_preempted(&mut output, kv_cache)?0;
    298
    299
            // 4. Schedule new sequences from waiting queue
    300
    12
            self.schedule_waiting(&mut output, kv_cache)?0;
    301
    302
            // 5. Build final output
    303
    25
            for &request_id13 in &self.running {
    304
    13
                if let Some(request) = self.requests.get(&request_id) {
    305
    13
                    if let Some(seq_id) = request.seq_id {
    306
    13
                        output.scheduled_seq_ids.push(seq_id);
    307
    13
                        output.scheduled_request_ids.push(request_id);
    308
    309
    13
                        if request.iterations == 0 {
    310
    11
                            // Prefill phase
    311
    11
                            output.num_prefill_tokens += request.input_ids.len();
    312
    11
                        } else {
    313
    2
                            // Decode phase (1 token per sequence)
    314
    2
                            output.num_decode_tokens += 1;
    315
    2
                        }
    316
    0
                    }
    317
    0
                }
    318
            }
    319
    320
    12
            self.stats.running_count = self.running.len();
    321
    12
            self.stats.queue_depth = self.waiting_queue.len();
    322
    323
    12
            Ok(output)
    324
    12
        }
    325
    326
        /// Update scheduler after generation iteration
    327
    4
        pub fn update_after_iteration(&mut self, generated_tokens: &HashMap<u64, u32>) {
    328
    8
            for (&request_id4, &token4) in generated_tokens {
    329
    4
                if let Some(request3) = self.requests.get_mut(&request_id) {
    330
    3
                    request.generated_tokens.push(token);
    331
    3
                    request.iterations += 1;
    332
    3
                }1
    333
            }
    334
    4
        }
    335
    336
        /// Mark request as complete
    337
    3
        pub fn complete_request(&mut self, request_id: u64, kv_cache: &mut PagedKvCache) {
    338
    3
            if let Some(request2) = self.requests.get_mut(&request_id) {
    339
    2
                request.state = SequenceState::Completed;
    340
    341
                // Free KV cache
    342
    2
                if let Some(seq_id) = request.seq_id {
    343
    2
                    kv_cache.free_sequence(seq_id);
    344
    2
                }0
    345
    346
                // Update stats
    347
    2
                self.stats.completed_requests += 1;
    348
    2
                let wait_time = request.wait_time().as_secs_f64() * 1000.0;
    349
    2
                self.total_wait_time_ms += wait_time;
    350
    2
                self.stats.avg_wait_time_ms =
    351
    2
                    self.total_wait_time_ms / self.stats.completed_requests as f64;
    352
    1
            }
    353
    354
            // Remove from running
    355
    3
            self.running.retain(|&id| id2 != request_id2);
    356
    3
        }
    357
    358
        /// Get request by ID
    359
    4
        pub fn get_request(&self, request_id: u64) -> Option<&SchedulerRequest> {
    360
    4
            self.requests.get(&request_id)
    361
    4
        }
    362
    363
        /// Get scheduler statistics
    364
    4
        pub fn stats(&self) -> &SchedulerStats {
    365
    4
            &self.stats
    366
    4
        }
    367
    368
        /// Check for completed sequences
    369
    12
        fn check_completions(&mut self, output: &mut SchedulerOutput, eos_token: u32) {
    370
    12
            let completed: Vec<u64> = self
    371
    12
                .running
    372
    12
                .iter()
    373
    12
                .filter(|&&id| {3
    374
    3
                    self.requests
    375
    3
                        .get(&id)
    376
    3
                        .is_some_and(|r| r.is_complete(eos_token))
    377
    3
                })
    378
    12
                .copied()
    379
    12
                .collect();
    380
    381
    13
            for request_id1 in completed {
    382
    1
                if let Some(request) = self.requests.get_mut(&request_id) {
    383
    1
                    request.state = SequenceState::Completed;
    384
    1
                }0
    385
    1
                output.completed_request_ids.push(request_id);
    386
            }
    387
    12
        }
    388
    389
        /// Handle preemption under memory pressure
    390
    12
        fn handle_preemption(&mut self, output: &mut SchedulerOutput, kv_cache: &mut PagedKvCache) {
    391
            // Simple preemption: if running at max and waiting queue has higher priority
    392
    12
            if self.running.len() >= self.max_batch_size && !self.waiting_queue.is_empty()1 {
    393
                // Check if waiting has higher priority than lowest running
    394
    1
                if let Some(waiting_entry) = self.waiting_queue.peek() {
    395
    1
                    let min_running_priority = self
    396
    1
                        .running
    397
    1
                        .iter()
    398
    1
                        .filter_map(|&id| self.requests.get(&id))
    399
    1
                        .map(|r| r.priority)
    400
    1
                        .min()
    401
    1
                        .unwrap_or(Priority::Critical);
    402
    403
    1
                    if waiting_entry.priority > min_running_priority {
    404
                        // Find lowest priority running request to preempt
    405
    1
                        if let Some(&preempt_id) = self.running.iter().find(|&&id| {
    406
    1
                            self.requests
    407
    1
                                .get(&id)
    408
    1
                                .is_some_and(|r| r.priority == min_running_priority)
    409
    1
                        }) {
    410
                            // Preempt the request
    411
    1
                            if let Some(request) = self.requests.get_mut(&preempt_id) {
    412
    1
                                request.state = SequenceState::Preempted;
    413
    1
                                if let Some(seq_id) = request.seq_id {
    414
    1
                                    output.preempted_seq_ids.push(seq_id);
    415
    1
                                    kv_cache.free_sequence(seq_id);
    416
    1
                                }0
    417
    1
                                request.seq_id = None;
    418
    0
                            }
    419
    1
                            self.running.retain(|&id| id != preempt_id);
    420
    1
                            self.preempted.push_back(preempt_id);
    421
    1
                            self.stats.preemptions += 1;
    422
    0
                        }
    423
    0
                    }
    424
    0
                }
    425
    11
            }
    426
    12
        }
    427
    428
        /// Resume preempted sequences
    429
    12
        fn resume_preempted(
    430
    12
            &mut self,
    431
    12
            _output: &mut SchedulerOutput,
    432
    12
            kv_cache: &mut PagedKvCache,
    433
    12
        ) -> Result<(), SchedulerError> {
    434
    13
            while self.running.len() < self.max_batch_size {
    435
    12
                if let Some(request_id1) = self.preempted.pop_front() {
    436
    1
                    if let Some(request) = self.requests.get_mut(&request_id) {
    437
                        // Try to allocate KV cache
    438
    1
                        let total_tokens = request.total_tokens();
    439
    1
                        match kv_cache.allocate_sequence(total_tokens) {
    440
    1
                            Ok(seq_id) => {
    441
    1
                                request.seq_id = Some(seq_id);
    442
    1
                                request.state = SequenceState::Running;
    443
    1
                                self.running.push(request_id);
    444
    1
                            },
    445
                            Err(_) => {
    446
                                // Can't allocate, put back in preempted queue
    447
    0
                                self.preempted.push_front(request_id);
    448
    0
                                break;
    449
                            },
    450
                        }
    451
    0
                    }
    452
                } else {
    453
    11
                    break;
    454
                }
    455
            }
    456
    12
            Ok(())
    457
    12
        }
    458
    459
        /// Schedule waiting requests
    460
    12
        fn schedule_waiting(
    461
    12
            &mut self,
    462
    12
            _output: &mut SchedulerOutput,
    463
    12
            kv_cache: &mut PagedKvCache,
    464
    12
        ) -> Result<(), SchedulerError> {
    465
    22
            while self.running.len() < self.max_batch_size {
    466
    18
                if let Some(entry10) = self.waiting_queue.pop() {
    467
    10
                    if let Some(request) = self.requests.get_mut(&entry.request_id) {
    468
    10
                        let total_tokens = request.input_ids.len();
    469
    10
                        match kv_cache.allocate_sequence(total_tokens) {
    470
    10
                            Ok(seq_id) => {
    471
    10
                                request.seq_id = Some(seq_id);
    472
    10
                                request.state = SequenceState::Running;
    473
    10
                                self.running.push(entry.request_id);
    474
    10
                            },
    475
                            Err(_) => {
    476
                                // Can't allocate, put back in queue (at front since already popped)
    477
    0
                                self.waiting_queue.push(entry);
    478
    0
                                break;
    479
                            },
    480
                        }
    481
    0
                    }
    482
                } else {
    483
    8
                    break;
    484
                }
    485
            }
    486
    12
            Ok(())
    487
    12
        }
    488
    489
        /// Number of waiting requests
    490
    4
        pub fn waiting_count(&self) -> usize {
    491
    4
            self.waiting_queue.len()
    492
    4
        }
    493
    494
        /// Number of running requests
    495
    4
        pub fn running_count(&self) -> usize {
    496
    4
            self.running.len()
    497
    4
        }
    498
    499
        /// Number of preempted requests
    500
    0
        pub fn preempted_count(&self) -> usize {
    501
    0
            self.preempted.len()
    502
    0
        }
    503
    }
    504
    505
    // ============================================================================
    506
    // SLOT-BASED SERVER CONCURRENCY (per llama.cpp)
    507
    // ============================================================================
    508
    //
    509
    // llama.cpp uses a slot-based architecture for concurrent inference:
    510
    // - Fixed number of slots, each with its own KV cache allocation
    511
    // - State machine: IDLE → PROCESSING → GENERATING → (complete) → IDLE
    512
    // - Slots can be dynamically assigned to incoming requests
    513
    // - Enables efficient handling of multiple concurrent clients
    514
    // ============================================================================
    515
    516
    /// Slot state machine (per llama.cpp server architecture)
    517
    ///
    518
    /// Each slot transitions through these states:
    519
    /// - IDLE: Ready to accept new requests
    520
    /// - PROCESSING: Initial prompt processing (prefill phase)
    521
    /// - GENERATING: Token generation (decode phase)
    522
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
    523
    pub enum SlotState {
    524
        /// Slot is available for new requests
    525
        Idle,
    526
        /// Processing initial prompt (prefill)
    527
        Processing,
    528
        /// Generating tokens (decode)
    529
        Generating,
    530
    }
    531
    532
    impl Default for SlotState {
    533
    1
        fn default() -> Self {
    534
    1
            Self::Idle
    535
    1
        }
    536
    }
    537
    538
    /// Server slot for concurrent inference
    539
    ///
    540
    /// Per llama.cpp: Each slot manages one inference request with its own
    541
    /// KV cache allocation and state machine.
    542
    #[derive(Debug, Clone)]
    543
    pub struct Slot {
    544
        /// Unique slot ID
    545
        pub id: usize,
    546
        /// Current state
    547
        pub state: SlotState,
    548
        /// Assigned request ID (None if idle)
    549
        pub request_id: Option<u64>,
    550
        /// Sequence ID for KV cache
    551
        pub seq_id: Option<SeqId>,
    552
        /// Input tokens (prompt)
    553
        pub input_tokens: Vec<u32>,
    554
        /// Generated tokens so far
    555
        pub generated_tokens: Vec<u32>,
    556
        /// Maximum tokens to generate
    557
        pub max_tokens: usize,
    558
        /// Number of prompt tokens processed
    559
        pub n_prompt_tokens_processed: usize,
    560
        /// Generation start time
    561
        pub generation_start: Option<Instant>,
    562
        /// Total prompt processing time (ms)
    563
        pub prompt_time_ms: f64,
    564
        /// Total generation time (ms)
    565
        pub generation_time_ms: f64,
    566
    }
    567
    568
    impl Slot {
    569
        /// Create a new idle slot
    570
    43
        pub fn new(id: usize) -> Self {
    571
    43
            Self {
    572
    43
                id,
    573
    43
                state: SlotState::Idle,
    574
    43
                request_id: None,
    575
    43
                seq_id: None,
    576
    43
                input_tokens: Vec::new(),
    577
    43
                generated_tokens: Vec::new(),
    578
    43
                max_tokens: 0,
    579
    43
                n_prompt_tokens_processed: 0,
    580
    43
                generation_start: None,
    581
    43
                prompt_time_ms: 0.0,
    582
    43
                generation_time_ms: 0.0,
    583
    43
            }
    584
    43
        }
    585
    586
        /// Check if slot is available
    587
    53
        pub fn is_idle(&self) -> bool {
    588
    53
            self.state == SlotState::Idle
    589
    53
        }
    590
    591
        /// Check if slot is actively generating
    592
    10
        pub fn is_generating(&self) -> bool {
    593
    10
            self.state == SlotState::Generating
    594
    10
        }
    595
    596
        /// Assign a request to this slot
    597
    18
        pub fn assign(&mut self, request_id: u64, input_tokens: Vec<u32>, max_tokens: usize) {
    598
    18
            self.state = SlotState::Processing;
    599
    18
            self.request_id = Some(request_id);
    600
    18
            self.input_tokens = input_tokens;
    601
    18
            self.max_tokens = max_tokens;
    602
    18
            self.generated_tokens.clear();
    603
    18
            self.n_prompt_tokens_processed = 0;
    604
    18
            self.prompt_time_ms = 0.0;
    605
    18
            self.generation_time_ms = 0.0;
    606
    18
            self.generation_start = None;
    607
    18
        }
    608
    609
        /// Transition from processing to generating
    610
    9
        pub fn start_generation(&mut self, prompt_time_ms: f64) {
    611
    9
            self.state = SlotState::Generating;
    612
    9
            self.prompt_time_ms = prompt_time_ms;
    613
    9
            self.generation_start = Some(Instant::now());
    614
    9
        }
    615
    616
        /// Add a generated token
    617
    8
        pub fn add_token(&mut self, token: u32) {
    618
    8
            self.generated_tokens.push(token);
    619
    8
        }
    620
    621
        /// Check if generation is complete
    622
    5
        pub fn is_complete(&self, eos_token: u32) -> bool {
    623
    5
            if self.generated_tokens.len() >= self.max_tokens {
    624
    1
                return true;
    625
    4
            }
    626
    4
            if let Some(&last3) = self.generated_tokens.last() {
    627
    3
                if last == eos_token {
    628
    1
                    return true;
    629
    2
                }
    630
    1
            }
    631
    3
            false
    632
    5
        }
    633
    634
        /// Finish generation and reset to idle
    635
    1
        pub fn finish(&mut self) {
    636
    1
            if let Some(start) = self.generation_start {
    637
    1
                self.generation_time_ms = start.elapsed().as_secs_f64() * 1000.0;
    638
    1
            }0
    639
    1
            self.state = SlotState::Idle;
    640
    1
            self.request_id = None;
    641
    1
            self.seq_id = None;
    642
    1
        }
    643
    644
        /// Get tokens per second for this slot's generation
    645
    5
        pub fn tokens_per_second(&self) -> f64 {
    646
    5
            if self.generation_time_ms > 0.0 {
    647
    0
                self.generated_tokens.len() as f64 / (self.generation_time_ms / 1000.0)
    648
            } else {
    649
    5
                0.0
    650
            }
    651
    5
        }
    652
    }
    653
    654
    /// Slot manager for concurrent inference
    655
    ///
    656
    /// Per llama.cpp: Manages a fixed pool of slots for handling concurrent requests.
    657
    /// Each slot has its own KV cache allocation and can process one request at a time.
    658
    #[derive(Debug)]
    659
    pub struct SlotManager {
    660
        /// Available slots
    661
        slots: Vec<Slot>,
    662
        /// Maximum context length per slot
    663
        pub max_context_length: usize,
    664
        /// Next request ID
    665
        next_request_id: u64,
    666
    }
    667
    668
    impl SlotManager {
    669
        /// Create a new slot manager with the specified number of slots
    670
    10
        pub fn new(num_slots: usize, max_context_length: usize) -> Self {
    671
    10
            let slots = (0..num_slots).map(Slot::new).collect();
    672
    10
            Self {
    673
    10
                slots,
    674
    10
                max_context_length,
    675
    10
                next_request_id: 0,
    676
    10
            }
    677
    10
        }
    678
    679
        /// Get number of total slots
    680
    1
        pub fn num_slots(&self) -> usize {
    681
    1
            self.slots.len()
    682
    1
        }
    683
    684
        /// Get number of idle slots
    685
    7
        pub fn num_idle_slots(&self) -> usize {
    686
    28
            self.slots.iter()7.filter7(|s| s.is_idle()).count7()
    687
    7
        }
    688
    689
        /// Get number of active (non-idle) slots
    690
    5
        pub fn num_active_slots(&self) -> usize {
    691
    5
            self.slots.len() - self.num_idle_slots()
    692
    5
        }
    693
    694
        /// Find an idle slot
    695
    12
        pub fn find_idle_slot(&self) -> Option<usize> {
    696
    12
            self.slots.iter().position(Slot::is_idle)
    697
    12
        }
    698
    699
        /// Assign a request to an available slot
    700
        ///
    701
        /// Returns the slot ID if successful, None if no slots available.
    702
    12
        pub fn assign_request(
    703
    12
            &mut self,
    704
    12
            input_tokens: Vec<u32>,
    705
    12
            max_tokens: usize,
    706
    12
        ) -> Option<(usize, u64)> {
    707
    12
            let slot_id11 = self.find_idle_slot()?1;
    708
    11
            let request_id = self.next_request_id;
    709
    11
            self.next_request_id += 1;
    710
    711
    11
            self.slots[slot_id].assign(request_id, input_tokens, max_tokens);
    712
    11
            Some((slot_id, request_id))
    713
    12
        }
    714
    715
        /// Get a reference to a slot
    716
    3
        pub fn get_slot(&self, slot_id: usize) -> Option<&Slot> {
    717
    3
            self.slots.get(slot_id)
    718
    3
        }
    719
    720
        /// Get a mutable reference to a slot
    721
    3
        pub fn get_slot_mut(&mut self, slot_id: usize) -> Option<&mut Slot> {
    722
    3
            self.slots.get_mut(slot_id)
    723
    3
        }
    724
    725
        /// Get all active slots (non-idle)
    726
    1
        pub fn active_slots(&self) -> impl Iterator<Item = &Slot> {
    727
    4
            self.slots.iter()1.filter1(|s| !s.is_idle())
    728
    1
        }
    729
    730
        /// Get all generating slots
    731
    1
        pub fn generating_slots(&self) -> impl Iterator<Item = &Slot> {
    732
    4
            self.slots.iter()1.filter1(|s| s.is_generating())
    733
    1
        }
    734
    735
        /// Get slots ready for batch processing
    736
    1
        pub fn batch_slots(&self) -> Vec<usize> {
    737
    1
            self.slots
    738
    1
                .iter()
    739
    1
                .enumerate()
    740
    4
                .filter1(|(_, s)| s.is_generating())
    741
    1
                .map(|(i, _)| i)
    742
    1
                .collect()
    743
    1
        }
    744
    745
        /// Get server utilization (0.0 to 1.0)
    746
    4
        pub fn utilization(&self) -> f64 {
    747
    4
            if self.slots.is_empty() {
    748
    1
                0.0
    749
            } else {
    750
    3
                self.num_active_slots() as f64 / self.slots.len() as f64
    751
            }
    752
    4
        }
    753
    754
        /// Get aggregate tokens per second across all slots
    755
    1
        pub fn aggregate_tokens_per_second(&self) -> f64 {
    756
    1
            self.slots.iter().map(Slot::tokens_per_second).sum()
    757
    1
        }
    758
    }
    759
    760
    // ============================================================================
    761
    // CONTINUOUS BATCHING (ubatch/sbatch per llama.cpp)
    762
    // ============================================================================
    763
    //
    764
    // llama.cpp's continuous batching system:
    765
    // - ubatch (micro-batch): Tokens processed in a single forward pass
    766
    // - sbatch (sequence batch): Multiple sequences grouped for batched inference
    767
    //
    768
    // This enables:
    769
    // - Dynamic batching: New sequences can join mid-inference
    770
    // - Efficient GPU utilization: Batch multiple decode steps together
    771
    // - Mixed prefill/decode: Process prefill and decode in same batch
    772
    // ============================================================================
    773
    774
    /// Batch type for continuous batching
    775
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
    776
    pub enum BatchType {
    777
        /// Prefill batch (processing initial prompts)
    778
        Prefill,
    779
        /// Decode batch (generating tokens)
    780
        Decode,
    781
        /// Mixed prefill and decode
    782
        Mixed,
    783
    }
    784
    785
    impl Default for BatchType {
    786
    1
        fn default() -> Self {
    787
    1
            Self::Decode
    788
    1
        }
    789
    }
    790
    791
    /// Token position within a batch
    792
    #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
    793
    pub struct BatchToken {
    794
        /// Token ID
    795
        pub token_id: u32,
    796
        /// Sequence ID this token belongs to
    797
        pub seq_idx: usize,
    798
        /// Position within the sequence
    799
        pub seq_pos: usize,
    800
        /// Whether this is a prompt token (vs generated)
    801
        pub is_prompt: bool,
    802
    }
    803
    804
    impl BatchToken {
    805
        /// Create a new batch token
    806
    32
        pub fn new(token_id: u32, seq_idx: usize, seq_pos: usize, is_prompt: bool) -> Self {
    807
    32
            Self {
    808
    32
                token_id,
    809
    32
                seq_idx,
    810
    32
                seq_pos,
    811
    32
                is_prompt,
    812
    32
            }
    813
    32
        }
    814
    }
    815
    816
    /// Micro-batch (ubatch) - tokens for a single forward pass
    817
    ///
    818
    /// Per llama.cpp: A ubatch contains tokens that will be processed together
    819
    /// in a single forward pass. Can contain tokens from multiple sequences.
    820
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    821
    pub struct MicroBatch {
    822
        /// Tokens in this micro-batch
    823
        pub tokens: Vec<BatchToken>,
    824
        /// Sequence indices included in this batch
    825
        pub seq_indices: Vec<usize>,
    826
        /// Batch type (prefill/decode/mixed)
    827
        pub batch_type: BatchType,
    828
        /// Maximum sequence length in batch
    829
        pub max_seq_len: usize,
    830
        /// Number of prompt tokens in batch
    831
        pub n_prompt_tokens: usize,
    832
        /// Number of decode tokens in batch
    833
        pub n_decode_tokens: usize,
    834
    }
    835
    836
    impl MicroBatch {
    837
        /// Create a new empty micro-batch
    838
    9
        pub fn new() -> Self {
    839
    9
            Self {
    840
    9
                tokens: Vec::new(),
    841
    9
                seq_indices: Vec::new(),
    842
    9
                batch_type: BatchType::Decode,
    843
    9
                max_seq_len: 0,
    844
    9
                n_prompt_tokens: 0,
    845
    9
                n_decode_tokens: 0,
    846
    9
            }
    847
    9
        }
    848
    849
        /// Create a micro-batch with capacity
    850
    7
        pub fn with_capacity(capacity: usize) -> Self {
    851
    7
            Self {
    852
    7
                tokens: Vec::with_capacity(capacity),
    853
    7
                seq_indices: Vec::new(),
    854
    7
                batch_type: BatchType::Decode,
    855
    7
                max_seq_len: 0,
    856
    7
                n_prompt_tokens: 0,
    857
    7
                n_decode_tokens: 0,
    858
    7
            }
    859
    7
        }
    860
    861
        /// Add a token to the batch
    862
    31
        pub fn add_token(&mut self, token: BatchToken) {
    863
    31
            if token.is_prompt {
    864
    23
                self.n_prompt_tokens += 1;
    865
    23
            } else {
    866
    8
                self.n_decode_tokens += 1;
    867
    8
            }
    868
    869
            // Track sequence index
    870
    31
            if !self.seq_indices.contains(&token.seq_idx) {
    871
    20
                self.seq_indices.push(token.seq_idx);
    872
    20
            }11
    873
    874
            // Update max sequence length
    875
    31
            self.max_seq_len = self.max_seq_len.max(token.seq_pos + 1);
    876
    877
    31
            self.tokens.push(token);
    878
    879
            // Update batch type
    880
    31
            self.update_batch_type();
    881
    31
        }
    882
    883
        /// Update batch type based on token composition
    884
    31
        fn update_batch_type(&mut self) {
    885
    31
            self.batch_type = match (self.n_prompt_tokens > 0, self.n_decode_tokens > 0) {
    886
    23
                (true, false) => BatchType::Prefill,
    887
    5
                (true, true) => BatchType::Mixed,
    888
                // Both (false, true) and (false, false) result in Decode
    889
    3
                (false, _) => BatchType::Decode,
    890
            };
    891
    31
        }
    892
    893
        /// Total number of tokens
    894
    29
        pub fn len(&self) -> usize {
    895
    29
            self.tokens.len()
    896
    29
        }
    897
    898
        /// Check if batch is empty
    899
    15
        pub fn is_empty(&self) -> bool {
    900
    15
            self.tokens.is_empty()
    901
    15
        }
    902
    903
        /// Number of sequences in batch
    904
    3
        pub fn num_sequences(&self) -> usize {
    905
    3
            self.seq_indices.len()
    906
    3
        }
    907
    908
        /// Check if batch is pure prefill
    909
    5
        pub fn is_prefill(&self) -> bool {
    910
    5
            self.batch_type == BatchType::Prefill
    911
    5
        }
    912
    913
        /// Check if batch is pure decode
    914
    3
        pub fn is_decode(&self) -> bool {
    915
    3
            self.batch_type == BatchType::Decode
    916
    3
        }
    917
    918
        /// Check if batch is mixed
    919
    3
        pub fn is_mixed(&self) -> bool {
    920
    3
            self.batch_type == BatchType::Mixed
    921
    3
        }
    922
    923
        /// Get token IDs as a vector (for model input)
    924
    2
        pub fn token_ids(&self) -> Vec<u32> {
    925
    2
            self.tokens.iter().map(|t| t.token_id).collect()
    926
    2
        }
    927
    928
        /// Get sequence positions (for position embeddings)
    929
    1
        pub fn positions(&self) -> Vec<usize> {
    930
    1
            self.tokens.iter().map(|t| t.seq_pos).collect()
    931
    1
        }
    932
    933
        /// Clear the batch
    934
    1
        pub fn clear(&mut self) {
    935
    1
            self.tokens.clear();
    936
    1
            self.seq_indices.clear();
    937
    1
            self.batch_type = BatchType::Decode;
    938
    1
            self.max_seq_len = 0;
    939
    1
            self.n_prompt_tokens = 0;
    940
    1
            self.n_decode_tokens = 0;
    941
    1
        }
    942
    }
    943
    944
    /// Sequence batch entry
    945
    #[derive(Debug, Clone, Serialize, Deserialize)]
    946
    pub struct SequenceBatchEntry {
    947
        /// Sequence index
    948
        pub seq_idx: usize,
    949
        /// Slot ID
    950
        pub slot_id: usize,
    951
        /// Request ID
    952
        pub request_id: u64,
    953
        /// Current position in sequence
    954
        pub position: usize,
    955
        /// Tokens to process (for prefill)
    956
        pub tokens: Vec<u32>,
    957
        /// Is this sequence in prefill or decode mode
    958
        pub is_prefill: bool,
    959
    }
    960
    961
    impl SequenceBatchEntry {
    962
        /// Create new sequence batch entry
    963
    28
        pub fn new(seq_idx: usize, slot_id: usize, request_id: u64) -> Self {
    964
    28
            Self {
    965
    28
                seq_idx,
    966
    28
                slot_id,
    967
    28
                request_id,
    968
    28
                position: 0,
    969
    28
                tokens: Vec::new(),
    970
    28
                is_prefill: true,
    971
    28
            }
    972
    28
        }
    973
    974
        /// Set tokens for prefill
    975
    12
        pub fn with_tokens(mut self, tokens: Vec<u32>) -> Self {
    976
    12
            self.tokens = tokens;
    977
    12
            self
    978
    12
        }
    979
    980
        /// Set position
    981
    1
        pub fn at_position(mut self, position: usize) -> Self {
    982
    1
            self.position = position;
    983
    1
            self
    984
    1
        }
    985
    986
        /// Mark as decode (not prefill)
    987
    4
        pub fn decoding(mut self) -> Self {
    988
    4
            self.is_prefill = false;
    989
    4
            self
    990
    4
        }
    991
    }
    992
    993
    /// Sequence batch (sbatch) - multiple sequences for batched inference
    994
    ///
    995
    /// Per llama.cpp: Groups sequences that will be processed together.
    996
    /// Manages the mapping from batch positions to individual sequences.
    997
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    998
    pub struct SequenceBatch {
    999
        /// Sequences in this batch
    1000
        pub sequences: Vec<SequenceBatchEntry>,
    1001
        /// Maximum batch size
    1002
        pub max_batch_size: usize,
    1003
        /// Current batch utilization
    1004
        pub utilization: f64,
    1005
    }
    1006
    1007
    impl SequenceBatch {
    1008
        /// Create a new sequence batch with max size
    1009
    25
        pub fn new(max_batch_size: usize) -> Self {
    1010
    25
            Self {
    1011
    25
                sequences: Vec::with_capacity(max_batch_size),
    1012
    25
                max_batch_size,
    1013
    25
                utilization: 0.0,
    1014
    25
            }
    1015
    25
        }
    1016
    1017
        /// Add a sequence to the batch
    1018
    26
        pub fn add_sequence(&mut self, entry: SequenceBatchEntry) -> bool {
    1019
    26
            if self.sequences.len() >= self.max_batch_size {
    1020
    1
                return false;
    1021
    25
            }
    1022
    25
            self.sequences.push(entry);
    1023
    25
            self.update_utilization();
    1024
    25
            true
    1025
    26
        }
    1026
    1027
        /// Remove a sequence by index
    1028
    4
        pub fn remove_sequence(&mut self, seq_idx: usize) -> Option<SequenceBatchEntry> {
    1029
    4
            let pos2 = self.sequences.iter().position(|s| s.seq_idx3 == seq_idx3)?2;
    1030
    2
            let entry = self.sequences.remove(pos);
    1031
    2
            self.update_utilization();
    1032
    2
            Some(entry)
    1033
    4
        }
    1034
    1035
        /// Update utilization metric
    1036
    28
        fn update_utilization(&mut self) {
    1037
    28
            self.utilization = if self.max_batch_size > 0 {
    1038
    27
                self.sequences.len() as f64 / self.max_batch_size as f64
    1039
            } else {
    1040
    1
                0.0
    1041
            };
    1042
    28
        }
    1043
    1044
        /// Number of sequences in batch
    1045
    6
        pub fn len(&self) -> usize {
    1046
    6
            self.sequences.len()
    1047
    6
        }
    1048
    1049
        /// Check if batch is empty
    1050
    3
        pub fn is_empty(&self) -> bool {
    1051
    3
            self.sequences.is_empty()
    1052
    3
        }
    1053
    1054
        /// Check if batch is full
    1055
    17
        pub fn is_full(&self) -> bool {
    1056
    17
            self.sequences.len() >= self.max_batch_size
    1057
    17
        }
    1058
    1059
        /// Get prefill sequences
    1060
    1
        pub fn prefill_sequences(&self) -> impl Iterator<Item = &SequenceBatchEntry> {
    1061
    1
            self.sequences.iter().filter(|s| s.is_prefill)
    1062
    1
        }
    1063
    1064
        /// Get decode sequences
    1065
    1
        pub fn decode_sequences(&self) -> impl Iterator<Item = &SequenceBatchEntry> {
    1066
    2
            self.sequences.iter()1.filter1(|s| !s.is_prefill)
    1067
    1
        }
    1068
    1069
        /// Count prefill sequences
    1070
    1
        pub fn num_prefill(&self) -> usize {
    1071
    1
            self.sequences.iter().filter(|s| s.is_prefill).count()
    1072
    1
        }
    1073
    1074
        /// Count decode sequences
    1075
    1
        pub fn num_decode(&self) -> usize {
    1076
    3
            self.sequences.iter()1.filter1(|s| !s.is_prefill).count1()
    1077
    1
        }
    1078
    1079
        /// Clear the batch
    1080
    1
        pub fn clear(&mut self) {
    1081
    1
            self.sequences.clear();
    1082
    1
            self.utilization = 0.0;
    1083
    1
        }
    1084
    1085
        /// Get sequence by index
    1086
    4
        pub fn get(&self, seq_idx: usize) -> Option<&SequenceBatchEntry> {
    1087
    4
            self.sequences.iter().find(|s| s.seq_idx3 == seq_idx3)
    1088
    4
        }
    1089
    1090
        /// Get mutable sequence by index
    1091
    5
        pub fn get_mut(&mut self, seq_idx: usize) -> Option<&mut SequenceBatchEntry> {
    1092
    5
            self.sequences.iter_mut().find(|s| s.seq_idx == seq_idx)
    1093
    5
        }
    1094
    }
    1095
    1096
    /// Batch configuration for continuous batching
    1097
    #[derive(Debug, Clone, Serialize, Deserialize)]
    1098
    pub struct BatchConfig {
    1099
        /// Maximum tokens per micro-batch
    1100
        pub max_ubatch_tokens: usize,
    1101
        /// Maximum sequences per sequence batch
    1102
        pub max_sbatch_sequences: usize,
    1103
        /// Prefer pure decode batches (vs mixed)
    1104
        pub prefer_pure_decode: bool,
    1105
        /// Maximum context length
    1106
        pub max_context_length: usize,
    1107
        /// Enable dynamic batching (add sequences mid-inference)
    1108
        pub dynamic_batching: bool,
    1109
    }
    1110
    1111
    impl Default for BatchConfig {
    1112
    16
        fn default() -> Self {
    1113
    16
            Self {
    1114
    16
                max_ubatch_tokens: 512,
    1115
    16
                max_sbatch_sequences: 8,
    1116
    16
                prefer_pure_decode: true,
    1117
    16
                max_context_length: 2048,
    1118
    16
                dynamic_batching: true,
    1119
    16
            }
    1120
    16
        }
    1121
    }
    1122
    1123
    impl BatchConfig {
    1124
        /// Create batch config with custom max tokens
    1125
    2
        pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
    1126
    2
            self.max_ubatch_tokens = max_tokens;
    1127
    2
            self
    1128
    2
        }
    1129
    1130
        /// Create batch config with custom max sequences
    1131
    2
        pub fn with_max_sequences(mut self, max_seqs: usize) -> Self {
    1132
    2
            self.max_sbatch_sequences = max_seqs;
    1133
    2
            self
    1134
    2
        }
    1135
    }
    1136
    1137
    /// Batch scheduler for continuous batching
    1138
    ///
    1139
    /// Coordinates micro-batch and sequence batch creation,
    1140
    /// implementing llama.cpp-style continuous batching.
    1141
    pub struct BatchScheduler {
    1142
        /// Configuration
    1143
        config: BatchConfig,
    1144
        /// Current sequence batch
    1145
        sbatch: SequenceBatch,
    1146
        /// Next sequence index
    1147
        next_seq_idx: usize,
    1148
        /// Statistics
    1149
        stats: BatchStats,
    1150
    }
    1151
    1152
    /// Statistics for batch scheduler
    1153
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    1154
    pub struct BatchStats {
    1155
        /// Total micro-batches created
    1156
        pub ubatches_created: u64,
    1157
        /// Total sequence batches created
    1158
        pub sbatches_created: u64,
    1159
        /// Total tokens processed
    1160
        pub tokens_processed: u64,
    1161
        /// Total prefill tokens
    1162
        pub prefill_tokens: u64,
    1163
        /// Total decode tokens
    1164
        pub decode_tokens: u64,
    1165
        /// Average tokens per ubatch
    1166
        pub avg_ubatch_size: f64,
    1167
        /// Average sequences per sbatch
    1168
        pub avg_sbatch_size: f64,
    1169
    }
    1170
    1171
    impl BatchScheduler {
    1172
        /// Create a new batch scheduler with default config
    1173
    11
        pub fn new() -> Self {
    1174
    11
            Self::with_config(BatchConfig::default())
    1175
    11
        }
    1176
    1177
        /// Create a new batch scheduler with custom config
    1178
    14
        pub fn with_config(config: BatchConfig) -> Self {
    1179
    14
            let max_seqs = config.max_sbatch_sequences;
    1180
    14
            Self {
    1181
    14
                config,
    1182
    14
                sbatch: SequenceBatch::new(max_seqs),
    1183
    14
                next_seq_idx: 0,
    1184
    14
                stats: BatchStats::default(),
    1185
    14
            }
    1186
    14
        }
    1187
    1188
        /// Add a new sequence to the batch scheduler
    1189
    12
        pub fn add_sequence(
    1190
    12
            &mut self,
    1191
    12
            slot_id: usize,
    1192
    12
            request_id: u64,
    1193
    12
            input_tokens: Vec<u32>,
    1194
    12
        ) -> Option<usize> {
    1195
    12
            if self.sbatch.is_full() {
    1196
    1
                return None;
    1197
    11
            }
    1198
    1199
    11
            let seq_idx = self.next_seq_idx;
    1200
    11
            self.next_seq_idx += 1;
    1201
    1202
    11
            let entry = SequenceBatchEntry::new(seq_idx, slot_id, request_id).with_tokens(input_tokens);
    1203
    1204
    11
            if self.sbatch.add_sequence(entry) {
    1205
    11
                Some(seq_idx)
    1206
            } else {
    1207
    0
                None
    1208
            }
    1209
    12
        }
    1210
    1211
        /// Mark sequence as completed and remove from batch
    1212
    2
        pub fn complete_sequence(&mut self, seq_idx: usize) -> Option<SequenceBatchEntry> {
    1213
    2
            self.sbatch.remove_sequence(seq_idx)
    1214
    2
        }
    1215
    1216
        /// Transition sequence from prefill to decode
    1217
    4
        pub fn start_decode(&mut self, seq_idx: usize, position: usize) -> bool {
    1218
    4
            if let Some(entry3) = self.sbatch.get_mut(seq_idx) {
    1219
    3
                entry.is_prefill = false;
    1220
    3
                entry.position = position;
    1221
    3
                entry.tokens.clear(); // No longer need prefill tokens
    1222
    3
                true
    1223
            } else {
    1224
    1
                false
    1225
            }
    1226
    4
        }
    1227
    1228
        /// Create a micro-batch from current sequences
    1229
        ///
    1230
        /// Returns a micro-batch optimized for the current state:
    1231
        /// - Prefill: Process all prompt tokens for prefill sequences
    1232
        /// - Decode: Process one token per decode sequence
    1233
        /// - Mixed: Combines both (if config allows)
    1234
    6
        pub fn create_ubatch(&mut self) -> MicroBatch {
    1235
    6
            let mut ubatch = MicroBatch::with_capacity(self.config.max_ubatch_tokens);
    1236
    1237
            // Process prefill sequences first (if any)
    1238
    12
            for entry6 in &self.sbatch.sequences {
    1239
    6
                if entry.is_prefill {
    1240
                    // Add all prefill tokens
    1241
    11
                    for (i, &token_id) in entry.tokens.iter()4.enumerate4() {
    1242
    11
                        if ubatch.len() >= self.config.max_ubatch_tokens {
    1243
    1
                            break;
    1244
    10
                        }
    1245
    10
                        ubatch.add_token(BatchToken::new(token_id, entry.seq_idx, i, true));
    1246
                    }
    1247
    2
                }
    1248
            }
    1249
    1250
            // If prefer_pure_decode and we have prefill tokens, return early
    1251
    6
            if self.config.prefer_pure_decode && !ubatch.is_empty()5 && ubatch3.is_prefill3() {
    1252
    3
                self.record_ubatch(&ubatch);
    1253
    3
                return ubatch;
    1254
    3
            }
    1255
    1256
            // Add decode tokens
    1257
    6
            for entry3 in &self.sbatch.sequences {
    1258
    3
                if !entry.is_prefill {
    1259
    2
                    if ubatch.len() >= self.config.max_ubatch_tokens {
    1260
    0
                        break;
    1261
    2
                    }
    1262
                    // Decode sequences process one token at their current position
    1263
                    // (the actual token ID will be filled in during inference)
    1264
    2
                    ubatch.add_token(BatchToken::new(
    1265
                        0, // Placeholder - will be filled by inference
    1266
    2
                        entry.seq_idx,
    1267
    2
                        entry.position,
    1268
                        false,
    1269
                    ));
    1270
    1
                }
    1271
            }
    1272
    1273
    3
            self.record_ubatch(&ubatch);
    1274
    3
            ubatch
    1275
    6
        }
    1276
    1277
        /// Record ubatch statistics
    1278
    6
        fn record_ubatch(&mut self, ubatch: &MicroBatch) {
    1279
    6
            if ubatch.is_empty() {
    1280
    1
                return;
    1281
    5
            }
    1282
    1283
    5
            self.stats.ubatches_created += 1;
    1284
    5
            self.stats.tokens_processed += ubatch.len() as u64;
    1285
    5
            self.stats.prefill_tokens += ubatch.n_prompt_tokens as u64;
    1286
    5
            self.stats.decode_tokens += ubatch.n_decode_tokens as u64;
    1287
    1288
            // Update rolling average
    1289
    5
            let n = self.stats.ubatches_created as f64;
    1290
    5
            self.stats.avg_ubatch_size =
    1291
    5
                self.stats.avg_ubatch_size * (n - 1.0) / n + ubatch.len() as f64 / n;
    1292
    6
        }
    1293
    1294
        /// Get the current sequence batch
    1295
    2
        pub fn sbatch(&self) -> &SequenceBatch {
    1296
    2
            &self.sbatch
    1297
    2
        }
    1298
    1299
        /// Get scheduler statistics
    1300
    1
        pub fn stats(&self) -> &BatchStats {
    1301
    1
            &self.stats
    1302
    1
        }
    1303
    1304
        /// Get configuration
    1305
    0
        pub fn config(&self) -> &BatchConfig {
    1306
    0
            &self.config
    1307
    0
        }
    1308
    1309
        /// Number of active sequences
    1310
    5
        pub fn num_sequences(&self) -> usize {
    1311
    5
            self.sbatch.len()
    1312
    5
        }
    1313
    1314
        /// Check if scheduler has capacity
    1315
    3
        pub fn has_capacity(&self) -> bool {
    1316
    3
            !self.sbatch.is_full()
    1317
    3
        }
    1318
    1319
        /// Current batch utilization
    1320
    1
        pub fn utilization(&self) -> f64 {
    1321
    1
            self.sbatch.utilization
    1322
    1
        }
    1323
    }
    1324
    1325
    impl Default for BatchScheduler {
    1326
    1
        fn default() -> Self {
    1327
    1
            Self::new()
    1328
    1
        }
    1329
    }
    1330
    1331
    // ============================================================================
    1332
    // DYNAMIC BATCH PRIORITY SCHEDULING
    1333
    // ============================================================================
    1334
    //
    1335
    // Advanced priority scheduling with:
    1336
    // - Age-based priority promotion (prevent starvation)
    1337
    // - Deadline-aware scheduling (SLA support)
    1338
    // - Priority-weighted token budgets
    1339
    // - Multi-level feedback queue (MLFQ) style scheduling
    1340
    // - Fair share allocation across priority levels
    1341
    //
    1342
    // Reference: Orca (Yu et al., 2022), vLLM priority scheduling
    1343
    // ============================================================================
    1344
    1345
    /// Deadline specification for a request
    1346
    #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
    1347
    pub struct Deadline {
    1348
        /// Target completion time (milliseconds from arrival)
    1349
        pub target_latency_ms: u64,
    1350
        /// Hard deadline (must complete by this time, else drop)
    1351
        pub hard_deadline_ms: Option<u64>,
    1352
        /// Soft SLA target (percentage of requests meeting target)
    1353
        pub sla_target: f64,
    1354
    }
    1355
    1356
    impl Default for Deadline {
    1357
    4
        fn default() -> Self {
    1358
    4
            Self {
    1359
    4
                target_latency_ms: 1000, // 1 second default
    1360
    4
                hard_deadline_ms: None,
    1361
    4
                sla_target: 0.99, // 99% SLA
    1362
    4
            }
    1363
    4
        }
    1364
    }
    1365
    1366
    impl Deadline {
    1367
        /// Create a deadline with target latency
    1368
    3
        pub fn with_target(target_ms: u64) -> Self {
    1369
    3
            Self {
    1370
    3
                target_latency_ms: target_ms,
    1371
    3
                ..Default::default()
    1372
    3
            }
    1373
    3
        }
    1374
    1375
        /// Create a strict deadline with hard cutoff
    1376
    1
        pub fn strict(target_ms: u64, hard_ms: u64) -> Self {
    1377
    1
            Self {
    1378
    1
                target_latency_ms: target_ms,
    1379
    1
                hard_deadline_ms: Some(hard_ms),
    1380
    1
                sla_target: 1.0,
    1381
    1
            }
    1382
    1
        }
    1383
    }
    1384
    1385
    /// Dynamic priority configuration
    1386
    #[derive(Debug, Clone, Serialize, Deserialize)]
    1387
    pub struct DynamicPriorityConfig {
    1388
        /// Enable age-based priority promotion
    1389
        pub enable_age_promotion: bool,
    1390
        /// Time (ms) before promoting to next priority level
    1391
        pub promotion_interval_ms: u64,
    1392
        /// Maximum priority level after promotion (prevent runaway)
    1393
        pub max_promoted_priority: Priority,
    1394
        /// Token budget per priority level (proportion of batch)
    1395
        pub priority_budgets: [f64; 4], // Low, Normal, High, Critical
    1396
        /// Enable deadline-aware scheduling
    1397
        pub enable_deadline_scheduling: bool,
    1398
        /// Urgency boost factor for approaching deadlines
    1399
        pub urgency_factor: f64,
    1400
        /// Minimum tokens to allocate per request
    1401
        pub min_tokens_per_request: usize,
    1402
        /// Enable fair share scheduling
    1403
        pub enable_fair_share: bool,
    1404
    }
    1405
    1406
    impl Default for DynamicPriorityConfig {
    1407
    19
        fn default() -> Self {
    1408
    19
            Self {
    1409
    19
                enable_age_promotion: true,
    1410
    19
                promotion_interval_ms: 5000, // Promote after 5 seconds
    1411
    19
                max_promoted_priority: Priority::High,
    1412
    19
                // Budget allocation: Low=5%, Normal=30%, High=40%, Critical=25%
    1413
    19
                priority_budgets: [0.05, 0.30, 0.40, 0.25],
    1414
    19
                enable_deadline_scheduling: true,
    1415
    19
                urgency_factor: 2.0,
    1416
    19
                min_tokens_per_request: 1,
    1417
    19
                enable_fair_share: true,
    1418
    19
            }
    1419
    19
        }
    1420
    }
    1421
    1422
    impl DynamicPriorityConfig {
    1423
        /// Create config with custom budgets
    1424
    1
        pub fn with_budgets(budgets: [f64; 4]) -> Self {
    1425
    1
            Self {
    1426
    1
                priority_budgets: budgets,
    1427
    1
                ..Default::default()
    1428
    1
            }
    1429
    1
        }
    1430
    1431
        /// Disable age promotion
    1432
    3
        pub fn no_promotion(mut self) -> Self {
    1433
    3
            self.enable_age_promotion = false;
    1434
    3
            self
    1435
    3
        }
    1436
    1437
        /// Set promotion interval
    1438
    1
        pub fn with_promotion_interval(mut self, ms: u64) -> Self {
    1439
    1
            self.promotion_interval_ms = ms;
    1440
    1
            self
    1441
    1
        }
    1442
    }
    1443
    1444
    /// Request entry with dynamic priority tracking
    1445
    #[derive(Debug, Clone)]
    1446
    pub struct DynamicRequest {
    1447
        /// Base request data
    1448
        pub request_id: u64,
    1449
        /// Input tokens
    1450
        pub input_ids: Vec<u32>,
    1451
        /// Maximum tokens to generate
    1452
        pub max_tokens: usize,
    1453
        /// Original priority (as submitted)
    1454
        pub original_priority: Priority,
    1455
        /// Effective priority (may be promoted)
    1456
        pub effective_priority: Priority,
    1457
        /// Arrival time
    1458
        pub arrival_time: Instant,
    1459
        /// Deadline specification
    1460
        pub deadline: Option<Deadline>,
    1461
        /// Number of times priority was promoted
    1462
        pub promotions: u32,
    1463
        /// Current state
    1464
        pub state: SequenceState,
    1465
        /// Generated tokens
    1466
        pub generated_tokens: Vec<u32>,
    1467
        /// Sequence ID (when scheduled)
    1468
        pub seq_id: Option<SeqId>,
    1469
        /// Time-to-first-token (if started)
    1470
        pub ttft_ms: Option<f64>,
    1471
    }
    1472
    1473
    impl DynamicRequest {
    1474
        /// Create a new dynamic request
    1475
    30
        pub fn new(request_id: u64, input_ids: Vec<u32>, max_tokens: usize) -> Self {
    1476
    30
            Self {
    1477
    30
                request_id,
    1478
    30
                input_ids,
    1479
    30
                max_tokens,
    1480
    30
                original_priority: Priority::Normal,
    1481
    30
                effective_priority: Priority::Normal,
    1482
    30
                arrival_time: Instant::now(),
    1483
    30
                deadline: None,
    1484
    30
                promotions: 0,
    1485
    30
                state: SequenceState::Waiting,
    1486
    30
                generated_tokens: Vec::new(),
    1487
    30
                seq_id: None,
    1488
    30
                ttft_ms: None,
    1489
    30
            }
    1490
    30
        }
    1491
    1492
        /// Set priority
    1493
    23
        pub fn with_priority(mut self, priority: Priority) -> Self {
    1494
    23
            self.original_priority = priority;
    1495
    23
            self.effective_priority = priority;
    1496
    23
            self
    1497
    23
        }
    1498
    1499
        /// Set deadline
    1500
    2
        pub fn with_deadline(mut self, deadline: Deadline) -> Self {
    1501
    2
            self.deadline = Some(deadline);
    1502
    2
            self
    1503
    2
        }
    1504
    1505
        /// Wait time since arrival
    1506
    20
        pub fn wait_time_ms(&self) -> u64 {
    1507
    20
            self.arrival_time.elapsed().as_millis() as u64
    1508
    20
        }
    1509
    1510
        /// Check if deadline is approaching (within 2x target latency)
    1511
    1
        pub fn is_urgent(&self) -> bool {
    1512
    1
            if let Some(deadline0) = &self.deadline {
    1513
    0
                let elapsed = self.wait_time_ms();
    1514
    0
                elapsed >= deadline.target_latency_ms / 2
    1515
            } else {
    1516
    1
                false
    1517
            }
    1518
    1
        }
    1519
    1520
        /// Check if hard deadline has passed
    1521
    13
        pub fn is_expired(&self) -> bool {
    1522
    13
            if let Some(deadline1) = &self.deadline {
    1523
    1
                if let Some(hard0) = deadline.hard_deadline_ms {
    1524
    0
                    return self.wait_time_ms() > hard;
    1525
    1
                }
    1526
    12
            }
    1527
    13
            false
    1528
    13
        }
    1529
    1530
        /// Calculate urgency score (0.0 to 1.0+)
    1531
        /// Higher score = more urgent
    1532
    2
        pub fn urgency_score(&self) -> f64 {
    1533
    2
            if let Some(deadline1) = &self.deadline {
    1534
    1
                let elapsed = self.wait_time_ms() as f64;
    1535
    1
                let target = deadline.target_latency_ms as f64;
    1536
    1
                if target > 0.0 {
    1537
    0
                    elapsed / target
    1538
                } else {
    1539
    1
                    0.0
    1540
                }
    1541
            } else {
    1542
    1
                0.0
    1543
            }
    1544
    2
        }
    1545
    1546
        /// Remaining tokens to generate
    1547
    11
        pub fn remaining_tokens(&self) -> usize {
    1548
    11
            self.max_tokens.saturating_sub(self.generated_tokens.len())
    1549
    11
        }
    1550
    1551
        /// Total tokens (input + generated)
    1552
    2
        pub fn total_tokens(&self) -> usize {
    1553
    2
            self.input_ids.len() + self.generated_tokens.len()
    1554
    2
        }
    1555
    }
    1556
    1557
    /// Statistics for dynamic priority scheduling
    1558
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    1559
    pub struct DynamicSchedulerStats {
    1560
        /// Total requests processed
    1561
        pub total_requests: u64,
    1562
        /// Requests completed
    1563
        pub completed_requests: u64,
    1564
        /// Requests that met SLA
    1565
        pub sla_met: u64,
    1566
        /// Requests that missed SLA
    1567
        pub sla_missed: u64,
    1568
        /// Requests dropped (hard deadline exceeded)
    1569
        pub dropped_requests: u64,
    1570
        /// Total priority promotions
    1571
        pub promotions: u64,
    1572
        /// Average time-to-first-token (ms)
    1573
        pub avg_ttft_ms: f64,
    1574
        /// p99 time-to-first-token (ms)
    1575
        pub p99_ttft_ms: f64,
    1576
        /// Tokens allocated per priority level
    1577
        pub tokens_by_priority: [u64; 4],
    1578
        /// Current queue depth per priority
    1579
        pub queue_depth_by_priority: [usize; 4],
    1580
    }
    1581
    1582
    /// Dynamic batch priority scheduler
    1583
    ///
    1584
    /// Implements advanced priority scheduling with:
    1585
    /// - Age-based priority promotion (MLFQ-style)
    1586
    /// - Deadline-aware scheduling for SLA compliance
    1587
    /// - Fair share token budget allocation
    1588
    /// - Urgency-based boosting for time-sensitive requests
    1589
    pub struct DynamicPriorityScheduler {
    1590
        /// Configuration
    1591
        config: DynamicPriorityConfig,
    1592
        /// All requests by ID
    1593
        requests: HashMap<u64, DynamicRequest>,
    1594
        /// Priority queues (one per level)
    1595
        priority_queues: [VecDeque<u64>; 4],
    1596
        /// Running requests
    1597
        running: Vec<u64>,
    1598
        /// Next request ID
    1599
        next_request_id: u64,
    1600
        /// Statistics
    1601
        stats: DynamicSchedulerStats,
    1602
        /// TTFT samples for percentile calculation
    1603
        ttft_samples: Vec<f64>,
    1604
        /// Total batch token budget
    1605
        batch_token_budget: usize,
    1606
    }
    1607
    1608
    impl DynamicPriorityScheduler {
    1609
        /// Create a new dynamic priority scheduler
    1610
    14
        pub fn new(batch_token_budget: usize) -> Self {
    1611
    14
            Self::with_config(batch_token_budget, DynamicPriorityConfig::default())
    1612
    14
        }
    1613
    1614
        /// Create with custom configuration
    1615
    17
        pub fn with_config(batch_token_budget: usize, config: DynamicPriorityConfig) -> Self {
    1616
    17
            Self {
    1617
    17
                config,
    1618
    17
                requests: HashMap::new(),
    1619
    17
                priority_queues: [
    1620
    17
                    VecDeque::new(),
    1621
    17
                    VecDeque::new(),
    1622
    17
                    VecDeque::new(),
    1623
    17
                    VecDeque::new(),
    1624
    17
                ],
    1625
    17
                running: Vec::new(),
    1626
    17
                next_request_id: 0,
    1627
    17
                stats: DynamicSchedulerStats::default(),
    1628
    17
                ttft_samples: Vec::new(),
    1629
    17
                batch_token_budget,
    1630
    17
            }
    1631
    17
        }
    1632
    1633
        /// Add a request with priority and optional deadline
    1634
    22
        pub fn add_request(
    1635
    22
            &mut self,
    1636
    22
            input_ids: Vec<u32>,
    1637
    22
            max_tokens: usize,
    1638
    22
            priority: Priority,
    1639
    22
            deadline: Option<Deadline>,
    1640
    22
        ) -> u64 {
    1641
    22
            let request_id = self.next_request_id;
    1642
    22
            self.next_request_id += 1;
    1643
    1644
    22
            let mut request =
    1645
    22
                DynamicRequest::new(request_id, input_ids, max_tokens).with_priority(priority);
    1646
    22
            if let Some(d1) = deadline {
    1647
    1
                request = request.with_deadline(d);
    1648
    21
            }
    1649
    1650
            // Add to appropriate priority queue
    1651
    22
            let queue_idx = priority as usize;
    1652
    22
            self.priority_queues[queue_idx].push_back(request_id);
    1653
    22
            self.requests.insert(request_id, request);
    1654
    1655
    22
            self.stats.total_requests += 1;
    1656
    22
            self.update_queue_depths();
    1657
    1658
    22
            request_id
    1659
    22
        }
    1660
    1661
        /// Add a simple request (Normal priority, no deadline)
    1662
    3
        pub fn add_simple_request(&mut self, input_ids: Vec<u32>, max_tokens: usize) -> u64 {
    1663
    3
            self.add_request(input_ids, max_tokens, Priority::Normal, None)
    1664
    3
        }
    1665
    1666
        /// Perform age-based priority promotion
    1667
    8
        pub fn promote_aged_requests(&mut self) {
    1668
    8
            if !self.config.enable_age_promotion {
    1669
    1
                return;
    1670
    7
            }
    1671
    1672
    7
            let promotion_threshold = self.config.promotion_interval_ms;
    1673
    7
            let max_priority = self.config.max_promoted_priority;
    1674
    1675
            // Check each queue except Critical (can't promote beyond Critical)
    1676
    28
            for queue_idx21 in 0..3 {
    1677
    21
                let current_priority = match queue_idx {
    1678
    7
                    0 => Priority::Low,
    1679
    7
                    1 => Priority::Normal,
    1680
    7
                    2 => Priority::High,
    1681
    0
                    _ => continue,
    1682
                };
    1683
    1684
                // Skip if current priority is already at max promoted level
    1685
    21
                if current_priority >= max_priority {
    1686
    7
                    continue;
    1687
    14
                }
    1688
    1689
                // Find requests to promote
    1690
    14
                let mut to_promote = Vec::new();
    1691
    14
                for &request_id9 in &self.priority_queues[queue_idx] {
    1692
    9
                    if let Some(request) = self.requests.get(&request_id) {
    1693
    9
                        let promotions_time = promotion_threshold * (request.promotions as u64 + 1);
    1694
    9
                        if request.wait_time_ms() >= promotions_time {
    1695
    0
                            to_promote.push(request_id);
    1696
    9
                        }
    1697
    0
                    }
    1698
                }
    1699
    1700
                // Promote requests
    1701
    14
                for request_id0 in to_promote {
    1702
    0
                    self.promote_request(request_id);
    1703
    0
                }
    1704
            }
    1705
    8
        }
    1706
    1707
        /// Promote a single request to next priority level
    1708
    0
        fn promote_request(&mut self, request_id: u64) {
    1709
    0
            if let Some(request) = self.requests.get_mut(&request_id) {
    1710
    0
                let current_idx = request.effective_priority as usize;
    1711
    0
                let max_idx = self.config.max_promoted_priority as usize;
    1712
    1713
    0
                if current_idx < max_idx {
    1714
                    // Remove from current queue
    1715
    0
                    self.priority_queues[current_idx].retain(|&id| id != request_id);
    1716
    1717
                    // Promote
    1718
    0
                    let new_priority = match current_idx + 1 {
    1719
    0
                        1 => Priority::Normal,
    1720
    0
                        2 => Priority::High,
    1721
    0
                        3 => Priority::Critical,
    1722
    0
                        _ => return,
    1723
                    };
    1724
    0
                    request.effective_priority = new_priority;
    1725
    0
                    request.promotions += 1;
    1726
    1727
                    // Add to new queue (at front since it's been waiting)
    1728
    0
                    self.priority_queues[current_idx + 1].push_front(request_id);
    1729
    0
                    self.stats.promotions += 1;
    1730
    0
                }
    1731
    0
            }
    1732
    0
        }
    1733
    1734
        /// Drop expired requests (hard deadline exceeded)
    1735
    7
        pub fn drop_expired(&mut self) -> Vec<u64> {
    1736
    7
            let mut dropped = Vec::new();
    1737
    1738
    35
            for queue28 in &mut self.priority_queues {
    1739
    28
                let mut to_remove = Vec::new();
    1740
    28
                for &request_id12 in queue.iter() {
    1741
    12
                    if let Some(request) = self.requests.get(&request_id) {
    1742
    12
                        if request.is_expired() {
    1743
    0
                            to_remove.push(request_id);
    1744
    12
                        }
    1745
    0
                    }
    1746
                }
    1747
    1748
    28
                for request_id0 in to_remove {
    1749
    0
                    queue.retain(|&id| id != request_id);
    1750
    0
                    if let Some(mut request) = self.requests.remove(&request_id) {
    1751
    0
                        request.state = SequenceState::Failed;
    1752
    0
                        dropped.push(request_id);
    1753
    0
                        self.stats.dropped_requests += 1;
    1754
    0
                    }
    1755
                }
    1756
            }
    1757
    1758
    7
            self.update_queue_depths();
    1759
    7
            dropped
    1760
    7
        }
    1761
    1762
        /// Schedule requests using dynamic priority and token budgets
    1763
        ///
    1764
        /// Returns (scheduled_request_ids, tokens_allocated_per_request)
    1765
    7
        pub fn schedule(&mut self, available_slots: usize) -> Vec<(u64, usize)> {
    1766
            // First, handle promotions and expirations
    1767
    7
            self.promote_aged_requests();
    1768
    7
            self.drop_expired();
    1769
    1770
    7
            let mut scheduled = Vec::new();
    1771
    7
            let mut remaining_budget = self.batch_token_budget;
    1772
    7
            let mut remaining_slots = available_slots;
    1773
    1774
            // Calculate token budgets per priority level
    1775
    7
            let budgets: [usize; 4] = if self.config.enable_fair_share {
    1776
    6
                self.config
    1777
    6
                    .priority_budgets
    1778
    24
                    .map6(|b| (b * self.batch_token_budget as f64) as usize)
    1779
            } else {
    1780
    1
                [
    1781
    1
                    remaining_budget,
    1782
    1
                    remaining_budget,
    1783
    1
                    remaining_budget,
    1784
    1
                    remaining_budget,
    1785
    1
                ]
    1786
            };
    1787
    1788
            // Schedule from highest priority to lowest
    1789
    22
            for queue_idx in (0..4)7.rev7() {
    1790
    22
                if remaining_slots == 0 || remaining_budget == 018 {
    1791
    5
                    break;
    1792
    17
                }
    1793
    1794
    17
                let queue = &mut self.priority_queues[queue_idx];
    1795
    17
                let mut priority_budget = budgets[queue_idx].min(remaining_budget);
    1796
    1797
                // Sort queue by urgency for deadline-aware scheduling
    1798
    17
                if self.config.enable_deadline_scheduling {
    1799
    17
                    let mut sorted: Vec<_> = queue.iter().copied().collect();
    1800
    17
                    sorted.sort_by(|&a, &b| {0
    1801
    0
                        let req_a = self.requests.get(&a);
    1802
    0
                        let req_b = self.requests.get(&b);
    1803
    0
                        match (req_a, req_b) {
    1804
    0
                            (Some(ra), Some(rb)) => rb
    1805
    0
                                .urgency_score()
    1806
    0
                                .partial_cmp(&ra.urgency_score())
    1807
    0
                                .unwrap_or(std::cmp::Ordering::Equal),
    1808
    0
                            _ => std::cmp::Ordering::Equal,
    1809
                        }
    1810
    0
                    });
    1811
    17
                    *queue = sorted.into_iter().collect();
    1812
    0
                }
    1813
    1814
                // Schedule requests from this priority level
    1815
    17
                let mut scheduled_from_queue = Vec::new();
    1816
    17
                for &request_id9 in queue.iter() {
    1817
    9
                    if remaining_slots == 0 || priority_budget < self.config.min_tokens_per_request {
    1818
    0
                        break;
    1819
    9
                    }
    1820
    1821
    9
                    if let Some(request) = self.requests.get(&request_id) {
    1822
                        // Calculate tokens to allocate
    1823
    9
                        let tokens_needed = request.remaining_tokens().max(1);
    1824
    9
                        let tokens_to_allocate = tokens_needed
    1825
    9
                            .min(priority_budget)
    1826
    9
                            .max(self.config.min_tokens_per_request);
    1827
    1828
    9
                        if tokens_to_allocate > 0 {
    1829
    9
                            scheduled.push((request_id, tokens_to_allocate));
    1830
    9
                            scheduled_from_queue.push(request_id);
    1831
    9
                            priority_budget = priority_budget.saturating_sub(tokens_to_allocate);
    1832
    9
                            remaining_budget = remaining_budget.saturating_sub(tokens_to_allocate);
    1833
    9
                            remaining_slots -= 1;
    1834
    9
    1835
    9
                            // Track tokens by priority
    1836
    9
                            self.stats.tokens_by_priority[queue_idx] += tokens_to_allocate as u64;
    1837
    9
                        }0
    1838
    0
                    }
    1839
                }
    1840
    1841
                // Remove scheduled requests from queue and update state
    1842
    26
                for request_id9 in scheduled_from_queue {
    1843
    9
                    queue.retain(|&id| id != request_id);
    1844
    9
                    if let Some(request) = self.requests.get_mut(&request_id) {
    1845
    9
                        request.state = SequenceState::Running;
    1846
    9
                        self.running.push(request_id);
    1847
    1848
                        // Record TTFT if first time running
    1849
    9
                        if request.ttft_ms.is_none() {
    1850
    9
                            let ttft = request.wait_time_ms() as f64;
    1851
    9
                            request.ttft_ms = Some(ttft);
    1852
    9
                            self.ttft_samples.push(ttft);
    1853
    9
                        }0
    1854
    0
                    }
    1855
                }
    1856
            }
    1857
    1858
    7
            self.update_queue_depths();
    1859
    7
            scheduled
    1860
    7
        }
    1861
    1862
        /// Complete a request and update statistics
    1863
    3
        pub fn complete_request(&mut self, request_id: u64) -> Option<DynamicRequest> {
    1864
            // Remove from running
    1865
    3
            self.running.retain(|&id| id2 != request_id2);
    1866
    1867
    3
            if let Some(mut request2) = self.requests.remove(&request_id) {
    1868
    2
                request.state = SequenceState::Completed;
    1869
    2
                self.stats.completed_requests += 1;
    1870
    1871
                // Check SLA compliance
    1872
    2
                if let Some(deadline1) = &request.deadline {
    1873
    1
                    let elapsed = request.wait_time_ms();
    1874
    1
                    if elapsed <= deadline.target_latency_ms {
    1875
    1
                        self.stats.sla_met += 1;
    1876
    1
                    } else {
    1877
    0
                        self.stats.sla_missed += 1;
    1878
    0
                    }
    1879
    1
                }
    1880
    1881
                // Update average TTFT
    1882
    2
                self.update_ttft_stats();
    1883
    1884
    2
                Some(request)
    1885
            } else {
    1886
    1
                None
    1887
            }
    1888
    3
        }
    1889
    1890
        /// Update TTFT statistics
    1891
    2
        fn update_ttft_stats(&mut self) {
    1892
    2
            if self.ttft_samples.is_empty() {
    1893
    0
                return;
    1894
    2
            }
    1895
    1896
            // Average
    1897
    2
            let sum: f64 = self.ttft_samples.iter().sum();
    1898
    2
            self.stats.avg_ttft_ms = sum / self.ttft_samples.len() as f64;
    1899
    1900
            // P99
    1901
    2
            let mut sorted = self.ttft_samples.clone();
    1902
    2
            sorted.sort_by(|a, b| a0.partial_cmp0(b0).unwrap_or0(std::cmp::Ordering::Equal0));
    1903
    2
            let p99_idx = ((sorted.len() as f64) * 0.99) as usize;
    1904
    2
            self.stats.p99_ttft_ms = sorted
    1905
    2
                .get(p99_idx.min(sorted.len() - 1))
    1906
    2
                .copied()
    1907
    2
                .unwrap_or(0.0);
    1908
    2
        }
    1909
    1910
        /// Update queue depth statistics
    1911
    36
        fn update_queue_depths(&mut self) {
    1912
    144
            for (i, queue) in self.priority_queues36.iter36().enumerate36() {
    1913
    144
                self.stats.queue_depth_by_priority[i] = queue.len();
    1914
    144
            }
    1915
    36
        }
    1916
    1917
        /// Get a request by ID
    1918
    2
        pub fn get_request(&self, request_id: u64) -> Option<&DynamicRequest> {
    1919
    2
            self.requests.get(&request_id)
    1920
    2
        }
    1921
    1922
        /// Get statistics
    1923
    5
        pub fn stats(&self) -> &DynamicSchedulerStats {
    1924
    5
            &self.stats
    1925
    5
        }
    1926
    1927
        /// Get configuration
    1928
    1
        pub fn config(&self) -> &DynamicPriorityConfig {
    1929
    1
            &self.config
    1930
    1
        }
    1931
    1932
        /// Total waiting requests
    1933
    4
        pub fn waiting_count(&self) -> usize {
    1934
    4
            self.priority_queues.iter().map(VecDeque::len).sum()
    1935
    4
        }
    1936
    1937
        /// Running requests
    1938
    3
        pub fn running_count(&self) -> usize {
    1939
    3
            self.running.len()
    1940
    3
        }
    1941
    1942
        /// SLA compliance rate (0.0 to 1.0)
    1943
    2
        pub fn sla_compliance_rate(&self) -> f64 {
    1944
    2
            let total = self.stats.sla_met + self.stats.sla_missed;
    1945
    2
            if total == 0 {
    1946
    1
                1.0
    1947
            } else {
    1948
    1
                self.stats.sla_met as f64 / total as f64
    1949
            }
    1950
    2
        }
    1951
    1952
        /// Get queue depth for a priority level
    1953
    8
        pub fn queue_depth(&self, priority: Priority) -> usize {
    1954
    8
            self.priority_queues[priority as usize].len()
    1955
    8
        }
    1956
    }
    1957
    1958
    // ============================================================================
    1959
    // Tests
    1960
    // ============================================================================
    1961
    1962
    // Tests extracted to tests.rs (PMAT-802)
    1963
    #[cfg(test)]
    1964
    #[path = "tests.rs"]
    1965
    mod scheduler_tests;