## 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/gguf/batch_scheduler.rs

    Line
    Count
    Source
    1
    //! Request Batching Infrastructure for GPU-Accelerated Inference
    2
    //!
    3
    //! Extracted from gguf_monolith.rs (PMAT-802) for vertical production partitioning.
    4
    //!
    5
    //! ## Contents
    6
    //!
    7
    //! - `BatchGenerationStats`: Statistics for batch generation capabilities
    8
    //! - `PendingRequest`, `RequestBatch`, `BatchRequestCollector`: Request batching
    9
    //! - `BatchingConfig`: Batching configuration
    10
    //! - `SlotState`, `ContinuousBatchScheduler`: Continuous batching
    11
    //! - `SpeculativeConfig`, `SpeculativeDecoder`: Speculative decoding
    12
    //! - `GpuBufferPool`, `GpuBufferPoolStats`: GPU buffer management
    13
    //! - `AsyncCommandQueue`, `CommandSlot`, `AsyncQueueStats`: Async command queue
    14
    //! - `PrefixCache`, `PrefixCacheEntry`, `PrefixCacheStats`: Prefix caching
    15
    //! - `MultiRequestState`, `MultiSchedulerRequest`, `SchedulingPolicy`, `MultiRequestScheduler`: Multi-request scheduling
    16
    //! - `ChunkedPrefillConfig`, `ChunkProgress`, `ChunkedPrefill`, `ChunkedPrefillStats`: Chunked prefill
    17
    //!
    18
    //! ## Feature Gate
    19
    //!
    20
    //! This entire module is gated behind `#[cfg(feature = "gpu")]`.
    21
    22
    // Note: This module is feature-gated in mod.rs with #[cfg(feature = "gpu")]
    23
    #![allow(clippy::many_single_char_names)]
    24
    #![allow(clippy::similar_names)]
    25
    26
    27
    use super::runtime::OwnedQuantizedKVCache;
    28
    29
    /// Statistics for batch generation configuration
    30
    #[derive(Debug, Clone)]
    31
    pub struct BatchGenerationStats {
    32
        /// Whether GPU cache is ready
    33
        pub gpu_cache_ready: bool,
    34
        /// Memory used by GPU cache in GB
    35
        pub cache_memory_gb: f64,
    36
        /// Number of transformer layers
    37
        pub num_layers: usize,
    38
        /// Hidden dimension
    39
        pub hidden_dim: usize,
    40
        /// FFN intermediate dimension
    41
        pub intermediate_dim: usize,
    42
        /// Recommended batch size for GPU efficiency
    43
        pub recommended_batch_size: usize,
    44
        /// Maximum batch size before memory pressure
    45
        pub max_batch_size: usize,
    46
    }
    47
    48
    // ============================================================================
    49
    // PARITY-023: Request Batching Infrastructure
    50
    // ============================================================================
    51
    52
    /// A pending request waiting to be batched (PARITY-023)
    53
    #[cfg(feature = "gpu")]
    54
    #[derive(Debug, Clone)]
    55
    pub struct PendingRequest {
    56
        /// Request ID for tracking
    57
        pub id: u64,
    58
        /// Prompt tokens
    59
        pub prompt: Vec<u32>,
    60
        /// Maximum tokens to generate
    61
        pub max_tokens: usize,
    62
        /// Temperature for sampling
    63
        pub temperature: f32,
    64
        /// Top-k sampling
    65
        pub top_k: usize,
    66
        /// Time when request was submitted
    67
        pub submitted_at: std::time::Instant,
    68
    }
    69
    70
    #[cfg(feature = "gpu")]
    71
    impl PendingRequest {
    72
        /// Create a new pending request
    73
    14
        pub fn new(
    74
    14
            id: u64,
    75
    14
            prompt: Vec<u32>,
    76
    14
            max_tokens: usize,
    77
    14
            temperature: f32,
    78
    14
            top_k: usize,
    79
    14
        ) -> Self {
    80
    14
            Self {
    81
    14
                id,
    82
    14
                prompt,
    83
    14
                max_tokens,
    84
    14
                temperature,
    85
    14
                top_k,
    86
    14
                submitted_at: std::time::Instant::now(),
    87
    14
            }
    88
    14
        }
    89
    90
        /// Time spent waiting in queue
    91
    8
        pub fn wait_time(&self) -> std::time::Duration {
    92
    8
            self.submitted_at.elapsed()
    93
    8
        }
    94
    }
    95
    96
    /// A batch of requests ready for processing (PARITY-023)
    97
    #[cfg(feature = "gpu")]
    98
    #[derive(Debug)]
    99
    pub struct RequestBatch {
    100
        /// Requests in this batch
    101
        pub requests: Vec<PendingRequest>,
    102
        /// When batch was formed
    103
        pub formed_at: std::time::Instant,
    104
    }
    105
    106
    #[cfg(feature = "gpu")]
    107
    impl RequestBatch {
    108
        /// Create batch from requests
    109
    2
        pub fn new(requests: Vec<PendingRequest>) -> Self {
    110
    2
            Self {
    111
    2
                requests,
    112
    2
                formed_at: std::time::Instant::now(),
    113
    2
            }
    114
    2
        }
    115
    116
        /// Number of requests in batch
    117
    4
        pub fn size(&self) -> usize {
    118
    4
            self.requests.len()
    119
    4
        }
    120
    121
        /// Extract prompts for batch processing
    122
    1
        pub fn prompts(&self) -> Vec<Vec<u32>> {
    123
    5
            self.requests.iter()1.map1(|r| r.prompt.clone()).collect1()
    124
    1
        }
    125
    126
        /// Average wait time for requests in this batch
    127
    1
        pub fn avg_wait_time(&self) -> std::time::Duration {
    128
    1
            if self.requests.is_empty() {
    129
    0
                return std::time::Duration::ZERO;
    130
    1
            }
    131
    1
            let total: std::time::Duration = self.requests.iter().map(PendingRequest::wait_time).sum();
    132
    1
            total / self.requests.len() as u32
    133
    1
        }
    134
    }
    135
    136
    /// Request batch collector with configurable thresholds (PARITY-023)
    137
    ///
    138
    /// Collects incoming requests and forms batches when:
    139
    /// - Batch size reaches `batch_threshold`, OR
    140
    /// - Wait time exceeds `timeout_ms`
    141
    ///
    142
    /// This enables efficient GPU utilization by batching multiple requests.
    143
    #[cfg(feature = "gpu")]
    144
    pub struct BatchRequestCollector {
    145
        /// Pending requests
    146
        pending: std::sync::Mutex<Vec<PendingRequest>>,
    147
        /// Next request ID
    148
        next_id: std::sync::atomic::AtomicU64,
    149
        /// Batch size threshold (32 = GPU GEMM threshold from IMP-600)
    150
        pub batch_threshold: usize,
    151
        /// Maximum wait time before forcing batch formation (ms)
    152
        pub timeout_ms: u64,
    153
        /// Maximum batch size (memory limit)
    154
        pub max_batch_size: usize,
    155
    }
    156
    157
    #[cfg(feature = "gpu")]
    158
    impl BatchRequestCollector {
    159
        /// Create new collector with default thresholds
    160
        ///
    161
        /// Default: batch_threshold=32, timeout_ms=50, max_batch_size=64
    162
    0
        pub fn new() -> Self {
    163
    0
            Self {
    164
    0
                pending: std::sync::Mutex::new(Vec::new()),
    165
    0
                next_id: std::sync::atomic::AtomicU64::new(0),
    166
    0
                batch_threshold: 32,
    167
    0
                timeout_ms: 50,
    168
    0
                max_batch_size: 64,
    169
    0
            }
    170
    0
        }
    171
    172
        /// Create collector with custom thresholds
    173
    2
        pub fn with_thresholds(batch_threshold: usize, timeout_ms: u64, max_batch_size: usize) -> Self {
    174
    2
            Self {
    175
    2
                pending: std::sync::Mutex::new(Vec::new()),
    176
    2
                next_id: std::sync::atomic::AtomicU64::new(0),
    177
    2
                batch_threshold,
    178
    2
                timeout_ms,
    179
    2
                max_batch_size,
    180
    2
            }
    181
    2
        }
    182
    183
        /// Submit a request to the collector
    184
        ///
    185
        /// Returns the request ID for tracking
    186
    8
        pub fn submit(
    187
    8
            &self,
    188
    8
            prompt: Vec<u32>,
    189
    8
            max_tokens: usize,
    190
    8
            temperature: f32,
    191
    8
            top_k: usize,
    192
    8
        ) -> u64 {
    193
    8
            let id = self
    194
    8
                .next_id
    195
    8
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    196
    8
            let request = PendingRequest::new(id, prompt, max_tokens, temperature, top_k);
    197
    198
    8
            let mut pending = self.pending.lock().expect("Mutex poisoned");
    199
    8
            pending.push(request);
    200
    201
    8
            id
    202
    8
        }
    203
    204
        /// Check if batch is ready to be formed
    205
    3
        pub fn is_batch_ready(&self) -> bool {
    206
    3
            let pending = self.pending.lock().expect("Mutex poisoned");
    207
    3
            if pending.is_empty() {
    208
    0
                return false;
    209
    3
            }
    210
    211
            // Batch ready if threshold reached
    212
    3
            if pending.len() >= self.batch_threshold {
    213
    1
                return true;
    214
    2
            }
    215
    216
            // Batch ready if oldest request has waited too long
    217
    2
            if let Some(oldest) = pending.first() {
    218
    2
                let wait_ms = oldest.wait_time().as_millis() as u64;
    219
    2
                if wait_ms >= self.timeout_ms {
    220
    0
                    return true;
    221
    2
                }
    222
    0
            }
    223
    224
    2
            false
    225
    3
        }
    226
    227
        /// Collect a batch of requests
    228
        ///
    229
        /// Returns None if no requests are pending or batch not ready
    230
    1
        pub fn collect_batch(&self) -> Option<RequestBatch> {
    231
    1
            let mut pending = self.pending.lock().expect("Mutex poisoned");
    232
    1
            if pending.is_empty() {
    233
    0
                return None;
    234
    1
            }
    235
    236
            // Check if batch is ready (threshold or timeout)
    237
    1
            let ready = pending.len() >= self.batch_threshold
    238
    0
                || pending
    239
    0
                    .first()
    240
    0
                    .is_some_and(|r| r.wait_time().as_millis() as u64 >= self.timeout_ms);
    241
    242
    1
            if !ready {
    243
    0
                return None;
    244
    1
            }
    245
    246
            // Take up to max_batch_size requests
    247
    1
            let batch_size = pending.len().min(self.max_batch_size);
    248
    1
            let requests: Vec<PendingRequest> = pending.drain(..batch_size).collect();
    249
    250
    1
            Some(RequestBatch::new(requests))
    251
    1
        }
    252
    253
        /// Force collect all pending requests as a batch
    254
    0
        pub fn flush(&self) -> Option<RequestBatch> {
    255
    0
            let mut pending = self.pending.lock().expect("Mutex poisoned");
    256
    0
            if pending.is_empty() {
    257
    0
                return None;
    258
    0
            }
    259
    260
    0
            let requests: Vec<PendingRequest> = pending.drain(..).collect();
    261
    0
            Some(RequestBatch::new(requests))
    262
    0
        }
    263
    264
        /// Number of pending requests
    265
    4
        pub fn pending_count(&self) -> usize {
    266
    4
            self.pending.lock().expect("Mutex poisoned").len()
    267
    4
        }
    268
    269
        /// Total requests submitted
    270
    1
        pub fn total_submitted(&self) -> u64 {
    271
    1
            self.next_id.load(std::sync::atomic::Ordering::Relaxed)
    272
    1
        }
    273
    }
    274
    275
    #[cfg(feature = "gpu")]
    276
    impl Default for BatchRequestCollector {
    277
    0
        fn default() -> Self {
    278
    0
            Self::new()
    279
    0
        }
    280
    }
    281
    282
    /// Batching configuration for request collector (PARITY-023)
    283
    #[cfg(feature = "gpu")]
    284
    #[derive(Debug, Clone)]
    285
    pub struct BatchingConfig {
    286
        /// Minimum batch size to trigger GPU processing (32 from IMP-600)
    287
        pub batch_threshold: usize,
    288
        /// Maximum wait time before processing smaller batch (ms)
    289
        pub timeout_ms: u64,
    290
        /// Maximum batch size (memory limit)
    291
        pub max_batch_size: usize,
    292
        /// Whether to prefer latency (process immediately) or throughput (wait for batch)
    293
        pub prefer_throughput: bool,
    294
    }
    295
    296
    #[cfg(feature = "gpu")]
    297
    impl Default for BatchingConfig {
    298
    1
        fn default() -> Self {
    299
    1
            Self {
    300
    1
                batch_threshold: 32,
    301
    1
                timeout_ms: 50,
    302
    1
                max_batch_size: 64,
    303
    1
                prefer_throughput: true,
    304
    1
            }
    305
    1
        }
    306
    }
    307
    308
    #[cfg(feature = "gpu")]
    309
    impl BatchingConfig {
    310
        /// Config optimized for latency (smaller batches, shorter timeout)
    311
    1
        pub fn latency_optimized() -> Self {
    312
    1
            Self {
    313
    1
                batch_threshold: 8,
    314
    1
                timeout_ms: 10,
    315
    1
                max_batch_size: 32,
    316
    1
                prefer_throughput: false,
    317
    1
            }
    318
    1
        }
    319
    320
        /// Config optimized for throughput (larger batches, longer timeout)
    321
    1
        pub fn throughput_optimized() -> Self {
    322
    1
            Self {
    323
    1
                batch_threshold: 32,
    324
    1
                timeout_ms: 100,
    325
    1
                max_batch_size: 64,
    326
    1
                prefer_throughput: true,
    327
    1
            }
    328
    1
        }
    329
    }
    330
    331
    /// Slot state for continuous batching (PARITY-028)
    332
    #[cfg(feature = "gpu")]
    333
    #[derive(Debug, Clone)]
    334
    pub enum SlotState {
    335
        /// Slot is available for new request
    336
        Empty,
    337
        /// Slot has active request being generated
    338
        Active {
    339
            /// Unique request identifier
    340
            request_id: u64,
    341
            /// Input prompt tokens
    342
            prompt_tokens: Vec<u32>,
    343
            /// Tokens generated so far
    344
            generated_tokens: Vec<u32>,
    345
            /// Maximum tokens to generate
    346
            max_tokens: usize,
    347
            /// Sampling temperature
    348
            temperature: f32,
    349
            /// Top-k sampling parameter
    350
            top_k: usize,
    351
        },
    352
        /// Slot has completed request waiting for retrieval
    353
        Completed {
    354
            /// Unique request identifier
    355
            request_id: u64,
    356
            /// All generated tokens
    357
            generated_tokens: Vec<u32>,
    358
        },
    359
    }
    360
    361
    #[cfg(feature = "gpu")]
    362
    impl SlotState {
    363
        /// Check if slot is available
    364
    109
        pub fn is_empty(&self) -> bool {
    365
    109
            matches!33(self, Self::Empty)
    366
    109
        }
    367
    368
        /// Check if slot has active generation
    369
    123
        pub fn is_active(&self) -> bool {
    370
    123
            matches!101(self, Self::Active { .. })
    371
    123
        }
    372
    373
        /// Check if slot has completed request
    374
    3
        pub fn is_completed(&self) -> bool {
    375
    3
            matches!2(self, Self::Completed { .. })
    376
    3
        }
    377
    378
        /// Get request ID if slot has one
    379
    3
        pub fn request_id(&self) -> Option<u64> {
    380
    3
            match self {
    381
    1
                Self::Empty => None,
    382
    1
                Self::Active { request_id, .. } | Self::Completed { request_id, .. } => {
    383
    2
                    Some(*request_id)
    384
                },
    385
            }
    386
    3
        }
    387
    }
    388
    389
    /// Continuous batch scheduler (PARITY-028)
    390
    ///
    391
    /// Enables dynamic addition/removal of requests from a running batch:
    392
    /// - Requests are assigned to slots
    393
    /// - Each slot can be in Empty, Active, or Completed state
    394
    /// - New requests fill empty slots immediately
    395
    /// - Completed requests free their slots for reuse
    396
    ///
    397
    /// This maximizes GPU utilization by keeping the batch full.
    398
    #[cfg(feature = "gpu")]
    399
    pub struct ContinuousBatchScheduler {
    400
        /// Fixed-size array of slots
    401
        slots: std::sync::Mutex<Vec<SlotState>>,
    402
        /// KV caches for each slot (pre-allocated)
    403
        caches: std::sync::Mutex<Vec<OwnedQuantizedKVCache>>,
    404
        /// Total slots (max concurrent requests)
    405
        pub num_slots: usize,
    406
        /// Completed request IDs for polling
    407
        completed: std::sync::Mutex<Vec<(u64, Vec<u32>)>>,
    408
        /// Next request ID
    409
        next_id: std::sync::atomic::AtomicU64,
    410
    }
    411
    412
    #[cfg(feature = "gpu")]
    413
    impl ContinuousBatchScheduler {
    414
        /// Create scheduler with specified number of slots
    415
        ///
    416
        /// # Arguments
    417
        /// * `num_slots` - Maximum concurrent requests (typically 32-64)
    418
        /// * `num_layers` - Number of transformer layers (for KV cache)
    419
        /// * `hidden_dim` - Hidden dimension (for KV cache)
    420
        /// * `max_seq_len` - Maximum sequence length (for KV cache)
    421
    3
        pub fn new(num_slots: usize, num_layers: usize, hidden_dim: usize, max_seq_len: usize) -> Self {
    422
    3
            let slots = vec![SlotState::Empty; num_slots];
    423
    3
            let caches = (0..num_slots)
    424
    40
                .map3(|_| OwnedQuantizedKVCache::new(num_layers, hidden_dim, max_seq_len))
    425
    3
                .collect();
    426
    427
    3
            Self {
    428
    3
                slots: std::sync::Mutex::new(slots),
    429
    3
                caches: std::sync::Mutex::new(caches),
    430
    3
                num_slots,
    431
    3
                completed: std::sync::Mutex::new(Vec::new()),
    432
    3
                next_id: std::sync::atomic::AtomicU64::new(0),
    433
    3
            }
    434
    3
        }
    435
    436
        /// Submit a new request to the scheduler
    437
        ///
    438
        /// Returns request ID if slot available, None if all slots full
    439
    10
        pub fn submit(
    440
    10
            &self,
    441
    10
            prompt_tokens: Vec<u32>,
    442
    10
            max_tokens: usize,
    443
    10
            temperature: f32,
    444
    10
            top_k: usize,
    445
    10
        ) -> Option<u64> {
    446
    10
            let mut slots = self.slots.lock().expect("Mutex poisoned");
    447
    448
            // Find first empty slot
    449
    10
            let empty_idx9 = slots.iter().position(SlotState::is_empty)?1;
    450
    451
    9
            let request_id = self
    452
    9
                .next_id
    453
    9
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    454
    455
    9
            slots[empty_idx] = SlotState::Active {
    456
    9
                request_id,
    457
    9
                prompt_tokens,
    458
    9
                generated_tokens: Vec::new(),
    459
    9
                max_tokens,
    460
    9
                temperature,
    461
    9
                top_k,
    462
    9
            };
    463
    464
    9
            Some(request_id)
    465
    10
        }
    466
    467
        /// Get number of active slots
    468
    9
        pub fn active_count(&self) -> usize {
    469
    9
            let slots = self.slots.lock().expect("Mutex poisoned");
    470
    120
            slots.iter()9.filter9(|s| s.is_active()).count9()
    471
    9
        }
    472
    473
        /// Get number of empty slots
    474
    6
        pub fn empty_count(&self) -> usize {
    475
    6
            let slots = self.slots.lock().expect("Mutex poisoned");
    476
    80
            slots.iter()6.filter6(|s| s.is_empty()).count6()
    477
    6
        }
    478
    479
        /// Check if any slot has completed request
    480
    2
        pub fn has_completed(&self) -> bool {
    481
    2
            let completed = self.completed.lock().expect("Mutex poisoned");
    482
    2
            !completed.is_empty()
    483
    2
        }
    484
    485
        /// Retrieve completed request results
    486
    1
        pub fn poll_completed(&self) -> Vec<(u64, Vec<u32>)> {
    487
    1
            let mut completed = self.completed.lock().expect("Mutex poisoned");
    488
    1
            std::mem::take(&mut *completed)
    489
    1
        }
    490
    491
        /// Mark a request as completed and move to completed queue
    492
    1
        pub fn complete_request(&self, slot_idx: usize, tokens: Vec<u32>) {
    493
    1
            let mut slots = self.slots.lock().expect("Mutex poisoned");
    494
    1
            let mut completed = self.completed.lock().expect("Mutex poisoned");
    495
    496
    1
            if slot_idx < slots.len() {
    497
    1
                if let SlotState::Active { request_id, .. } = &slots[slot_idx] {
    498
    1
                    let id = *request_id;
    499
    1
                    // Move to completed
    500
    1
                    completed.push((id, tokens));
    501
    1
                    // Free the slot
    502
    1
                    slots[slot_idx] = SlotState::Empty;
    503
    1
    504
    1
                    // Reset KV cache for this slot
    505
    1
                    let mut caches = self.caches.lock().expect("Mutex poisoned");
    506
    1
                    caches[slot_idx].reset();
    507
    1
                }0
    508
    0
            }
    509
    1
        }
    510
    511
        /// Get active slot indices and their current positions
    512
    0
        pub fn get_active_slots(&self) -> Vec<(usize, usize)> {
    513
    0
            let slots = self.slots.lock().expect("Mutex poisoned");
    514
    0
            slots
    515
    0
                .iter()
    516
    0
                .enumerate()
    517
    0
                .filter_map(|(idx, slot)| match slot {
    518
                    SlotState::Active {
    519
    0
                        prompt_tokens,
    520
    0
                        generated_tokens,
    521
                        ..
    522
                    } => {
    523
    0
                        let pos = prompt_tokens.len() + generated_tokens.len();
    524
    0
                        Some((idx, pos))
    525
                    },
    526
    0
                    _ => None,
    527
    0
                })
    528
    0
                .collect()
    529
    0
        }
    530
    531
        /// Get utilization (active_slots / total_slots)
    532
    3
        pub fn utilization(&self) -> f64 {
    533
    3
            let active = self.active_count();
    534
    3
            active as f64 / self.num_slots as f64
    535
    3
        }
    536
    }
    537
    538
    /// Speculative decoding configuration (PARITY-029)
    539
    #[cfg(feature = "gpu")]
    540
    #[derive(Debug, Clone)]
    541
    pub struct SpeculativeConfig {
    542
        /// Number of tokens to speculatively generate per step
    543
        pub speculation_length: usize,
    544
        /// Temperature for draft model (lower = more deterministic)
    545
        pub draft_temperature: f32,
    546
        /// Whether to use same model for draft (self-speculative)
    547
        pub self_speculative: bool,
    548
    }
    549
    550
    #[cfg(feature = "gpu")]
    551
    impl Default for SpeculativeConfig {
    552
    4
        fn default() -> Self {
    553
    4
            Self {
    554
    4
                speculation_length: 4,
    555
    4
                draft_temperature: 0.0,
    556
    4
                self_speculative: true,
    557
    4
            }
    558
    4
        }
    559
    }
    560
    561
    /// Result of speculative decoding verification step
    562
    #[cfg(feature = "gpu")]
    563
    #[derive(Debug, Clone)]
    564
    pub struct VerificationResult {
    565
        /// Number of draft tokens accepted
    566
        pub accepted_count: usize,
    567
        /// Total draft tokens generated
    568
        pub draft_count: usize,
    569
        /// Accepted tokens (verified by target model)
    570
        pub accepted_tokens: Vec<u32>,
    571
        /// Whether all draft tokens were accepted
    572
        pub all_accepted: bool,
    573
    }
    574
    575
    /// Speculative decoder for accelerated token generation (PARITY-029)
    576
    ///
    577
    /// Implements speculative decoding (Leviathan et al., 2023):
    578
    /// 1. Draft model generates K candidate tokens quickly
    579
    /// 2. Target model verifies all K tokens in parallel
    580
    /// 3. Accept tokens until first rejection, then resample
    581
    ///
    582
    /// This enables O(K) speedup when draft acceptance rate is high.
    583
    #[cfg(feature = "gpu")]
    584
    pub struct SpeculativeDecoder {
    585
        /// Speculative decoding configuration
    586
        pub config: SpeculativeConfig,
    587
        /// Statistics: total draft tokens generated
    588
        pub total_draft_tokens: std::sync::atomic::AtomicU64,
    589
        /// Statistics: total draft tokens accepted
    590
        pub total_accepted_tokens: std::sync::atomic::AtomicU64,
    591
    }
    592
    593
    #[cfg(feature = "gpu")]
    594
    impl SpeculativeDecoder {
    595
        /// Create new speculative decoder with default config
    596
    2
        pub fn new() -> Self {
    597
    2
            Self {
    598
    2
                config: SpeculativeConfig::default(),
    599
    2
                total_draft_tokens: std::sync::atomic::AtomicU64::new(0),
    600
    2
                total_accepted_tokens: std::sync::atomic::AtomicU64::new(0),
    601
    2
            }
    602
    2
        }
    603
    604
        /// Create speculative decoder with custom config
    605
    1
        pub fn with_config(config: SpeculativeConfig) -> Self {
    606
    1
            Self {
    607
    1
                config,
    608
    1
                total_draft_tokens: std::sync::atomic::AtomicU64::new(0),
    609
    1
                total_accepted_tokens: std::sync::atomic::AtomicU64::new(0),
    610
    1
            }
    611
    1
        }
    612
    613
        /// Get acceptance rate (accepted / total draft tokens)
    614
    6
        pub fn acceptance_rate(&self) -> f64 {
    615
    6
            let total = self
    616
    6
                .total_draft_tokens
    617
    6
                .load(std::sync::atomic::Ordering::Relaxed);
    618
    6
            let accepted = self
    619
    6
                .total_accepted_tokens
    620
    6
                .load(std::sync::atomic::Ordering::Relaxed);
    621
    6
            if total == 0 {
    622
    4
                return 0.0;
    623
    2
            }
    624
    2
            accepted as f64 / total as f64
    625
    6
        }
    626
    627
        /// Verify draft tokens against target model logits
    628
        ///
    629
        /// # Arguments
    630
        /// * `draft_tokens` - Candidate tokens from draft model
    631
        /// * `target_logits` - Logits from target model for each position
    632
        /// * `temperature` - Sampling temperature for rejection sampling
    633
        ///
    634
        /// # Returns
    635
        /// VerificationResult with accepted tokens and statistics
    636
    12
        pub fn verify_draft(
    637
    12
            &self,
    638
    12
            draft_tokens: &[u32],
    639
    12
            target_logits: &[Vec<f32>],
    640
    12
            temperature: f32,
    641
    12
        ) -> VerificationResult {
    642
    12
            let mut accepted_tokens = Vec::with_capacity(draft_tokens.len());
    643
    12
            let mut accepted_count = 0;
    644
    645
            // Verify each draft token against target model distribution
    646
    44
            for (i, &draft_token) in draft_tokens12.iter12().enumerate12() {
    647
    44
                if i >= target_logits.len() {
    648
    0
                    break;
    649
    44
                }
    650
    651
    44
                let logits = &target_logits[i];
    652
    653
                // Find target model's top token
    654
    44
                let (target_token, _) = logits
    655
    44
                    .iter()
    656
    44
                    .enumerate()
    657
    4.35k
                    .max_by44(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    658
    44
                    .unwrap_or((0, &0.0));
    659
    660
                // Accept if draft matches target (greedy case)
    661
    44
                if temperature == 0.0 {
    662
    44
                    if draft_token == target_token as u32 {
    663
    40
                        accepted_tokens.push(draft_token);
    664
    40
                        accepted_count += 1;
    665
    40
                    } else {
    666
                        // Reject and use target's token instead
    667
    4
                        accepted_tokens.push(target_token as u32);
    668
    4
                        accepted_count += 1;
    669
    4
                        break; // Stop at first mismatch
    670
                    }
    671
                } else {
    672
                    // Rejection sampling for non-greedy decoding
    673
                    // P(accept) = min(1, p_target(x) / p_draft(x))
    674
                    // For simplicity, accept if draft is in top-k of target
    675
    0
                    let mut sorted_indices: Vec<usize> = (0..logits.len()).collect();
    676
    0
                    sorted_indices.sort_by(|&a, &b| {
    677
    0
                        logits[b]
    678
    0
                            .partial_cmp(&logits[a])
    679
    0
                            .unwrap_or(std::cmp::Ordering::Equal)
    680
    0
                    });
    681
    682
    0
                    let top_k = 10; // Accept if in top-10
    683
    0
                    let in_top_k = sorted_indices
    684
    0
                        .iter()
    685
    0
                        .take(top_k)
    686
    0
                        .any(|&idx| idx == draft_token as usize);
    687
    688
    0
                    if in_top_k {
    689
    0
                        accepted_tokens.push(draft_token);
    690
    0
                        accepted_count += 1;
    691
    0
                    } else {
    692
                        // Reject, use target's sampled token
    693
    0
                        accepted_tokens.push(sorted_indices[0] as u32);
    694
    0
                        accepted_count += 1;
    695
    0
                        break;
    696
                    }
    697
                }
    698
            }
    699
    700
            // Update statistics
    701
    12
            self.total_draft_tokens.fetch_add(
    702
    12
                draft_tokens.len() as u64,
    703
    12
                std::sync::atomic::Ordering::Relaxed,
    704
            );
    705
    12
            self.total_accepted_tokens
    706
    12
                .fetch_add(accepted_count as u64, std::sync::atomic::Ordering::Relaxed);
    707
    708
    12
            VerificationResult {
    709
    12
                accepted_count,
    710
    12
                draft_count: draft_tokens.len(),
    711
    12
                accepted_tokens,
    712
    12
                all_accepted: accepted_count == draft_tokens.len(),
    713
    12
            }
    714
    12
        }
    715
    716
        /// Calculate expected speedup based on acceptance rate
    717
        ///
    718
        /// Speedup = K * acceptance_rate + 1 (always get at least 1 token)
    719
    3
        pub fn expected_speedup(&self) -> f64 {
    720
    3
            let k = self.config.speculation_length as f64;
    721
    3
            let acceptance_rate = self.acceptance_rate();
    722
    3
            k * acceptance_rate + 1.0
    723
    3
        }
    724
    725
        /// Reset statistics
    726
    1
        pub fn reset_stats(&self) {
    727
    1
            self.total_draft_tokens
    728
    1
                .store(0, std::sync::atomic::Ordering::Relaxed);
    729
    1
            self.total_accepted_tokens
    730
    1
                .store(0, std::sync::atomic::Ordering::Relaxed);
    731
    1
        }
    732
    }
    733
    734
    #[cfg(feature = "gpu")]
    735
    impl Default for SpeculativeDecoder {
    736
    0
        fn default() -> Self {
    737
    0
            Self::new()
    738
    0
        }
    739
    }
    740
    741
    /// GPU Buffer Pool for zero-allocation inference (PARITY-031, IMP-309)
    742
    ///
    743
    /// Pre-allocates GPU buffers during warmup to eliminate allocation overhead
    744
    /// during generation. Uses a pool of reusable buffers for each tensor type.
    745
    ///
    746
    /// # Key Properties
    747
    /// - Zero GPU malloc after warmup phase
    748
    /// - Pre-allocated buffers for common tensor sizes
    749
    /// - Thread-safe buffer borrowing and return
    750
    ///
    751
    /// # Buffer Types
    752
    /// - Hidden state buffers: [batch_size, hidden_dim]
    753
    /// - Intermediate buffers: [batch_size, intermediate_dim]
    754
    /// - Attention score buffers: [batch_size, num_heads, seq_len]
    755
    /// - KV cache buffers: [num_layers, seq_len, hidden_dim]
    756
    #[cfg(feature = "gpu")]
    757
    pub struct GpuBufferPool {
    758
        /// Pre-allocated hidden state buffers
    759
        hidden_buffers: std::sync::Mutex<Vec<Vec<f32>>>,
    760
        /// Pre-allocated intermediate buffers (FFN)
    761
        intermediate_buffers: std::sync::Mutex<Vec<Vec<f32>>>,
    762
        /// Pre-allocated attention score buffers
    763
        attention_buffers: std::sync::Mutex<Vec<Vec<f32>>>,
    764
        /// Buffer dimensions for validation
    765
        hidden_dim: usize,
    766
        intermediate_dim: usize,
    767
        max_seq_len: usize,
    768
        num_heads: usize,
    769
        /// Pool size per buffer type
    770
        pool_size: usize,
    771
        /// Statistics: buffers borrowed
    772
        pub borrows: std::sync::atomic::AtomicU64,
    773
        /// Statistics: buffers returned
    774
        pub returns: std::sync::atomic::AtomicU64,
    775
        /// Statistics: allocations after warmup (should be 0)
    776
        pub post_warmup_allocs: std::sync::atomic::AtomicU64,
    777
        /// Whether warmup is complete
    778
        warmed_up: std::sync::atomic::AtomicBool,
    779
    }
    780
    781
    #[cfg(feature = "gpu")]
    782
    impl GpuBufferPool {
    783
        /// Create new buffer pool with specified dimensions
    784
    5
        pub fn new(
    785
    5
            hidden_dim: usize,
    786
    5
            intermediate_dim: usize,
    787
    5
            max_seq_len: usize,
    788
    5
            num_heads: usize,
    789
    5
            pool_size: usize,
    790
    5
        ) -> Self {
    791
    5
            Self {
    792
    5
                hidden_buffers: std::sync::Mutex::new(Vec::with_capacity(pool_size)),
    793
    5
                intermediate_buffers: std::sync::Mutex::new(Vec::with_capacity(pool_size)),
    794
    5
                attention_buffers: std::sync::Mutex::new(Vec::with_capacity(pool_size)),
    795
    5
                hidden_dim,
    796
    5
                intermediate_dim,
    797
    5
                max_seq_len,
    798
    5
                num_heads,
    799
    5
                pool_size,
    800
    5
                borrows: std::sync::atomic::AtomicU64::new(0),
    801
    5
                returns: std::sync::atomic::AtomicU64::new(0),
    802
    5
                post_warmup_allocs: std::sync::atomic::AtomicU64::new(0),
    803
    5
                warmed_up: std::sync::atomic::AtomicBool::new(false),
    804
    5
            }
    805
    5
        }
    806
    807
        /// Warmup: pre-allocate all buffers
    808
        ///
    809
        /// Call this once during model initialization to eliminate
    810
        /// allocation overhead during inference.
    811
    3
        pub fn warmup(&self) {
    812
            // Pre-allocate hidden state buffers
    813
            {
    814
    3
                let mut buffers = self.hidden_buffers.lock().expect("mutex poisoned");
    815
    16
                for _ in 0..self.pool_size3 {
    816
    16
                    buffers.push(vec![0.0f32; self.hidden_dim]);
    817
    16
                }
    818
            }
    819
    820
            // Pre-allocate intermediate buffers (FFN)
    821
            {
    822
    3
                let mut buffers = self.intermediate_buffers.lock().expect("mutex poisoned");
    823
    16
                for _ in 0..self.pool_size3 {
    824
    16
                    buffers.push(vec![0.0f32; self.intermediate_dim]);
    825
    16
                }
    826
            }
    827
    828
            // Pre-allocate attention score buffers
    829
            {
    830
    3
                let mut buffers = self.attention_buffers.lock().expect("mutex poisoned");
    831
    16
                for _ in 0..self.pool_size3 {
    832
    16
                    buffers.push(vec![0.0f32; self.num_heads * self.max_seq_len]);
    833
    16
                }
    834
            }
    835
    836
    3
            self.warmed_up
    837
    3
                .store(true, std::sync::atomic::Ordering::Release);
    838
    3
        }
    839
    840
        /// Borrow a hidden state buffer from the pool
    841
        ///
    842
        /// Returns a pre-allocated buffer if available, or allocates new if needed.
    843
    11
        pub fn borrow_hidden(&self) -> Vec<f32> {
    844
    11
            self.borrows
    845
    11
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    846
    847
    11
            let mut buffers = self.hidden_buffers.lock().expect("mutex poisoned");
    848
    11
            if let Some(buffer) = buffers.pop() {
    849
    11
                buffer
    850
            } else {
    851
                // Need to allocate - track if after warmup
    852
    0
                if self.warmed_up.load(std::sync::atomic::Ordering::Acquire) {
    853
    0
                    self.post_warmup_allocs
    854
    0
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    855
    0
                }
    856
    0
                vec![0.0f32; self.hidden_dim]
    857
            }
    858
    11
        }
    859
    860
        /// Return a hidden state buffer to the pool
    861
    11
        pub fn return_hidden(&self, mut buffer: Vec<f32>) {
    862
    11
            self.returns
    863
    11
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    864
    865
            // Zero out for security and determinism
    866
    11
            buffer.fill(0.0);
    867
    868
    11
            let mut buffers = self.hidden_buffers.lock().expect("mutex poisoned");
    869
    11
            if buffers.len() < self.pool_size {
    870
    11
                buffers.push(buffer);
    871
    11
            }0
    872
            // If pool is full, buffer is dropped
    873
    11
        }
    874
    875
        /// Borrow an intermediate buffer from the pool
    876
    10
        pub fn borrow_intermediate(&self) -> Vec<f32> {
    877
    10
            self.borrows
    878
    10
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    879
    880
    10
            let mut buffers = self.intermediate_buffers.lock().expect("mutex poisoned");
    881
    10
            if let Some(buffer) = buffers.pop() {
    882
    10
                buffer
    883
            } else {
    884
    0
                if self.warmed_up.load(std::sync::atomic::Ordering::Acquire) {
    885
    0
                    self.post_warmup_allocs
    886
    0
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    887
    0
                }
    888
    0
                vec![0.0f32; self.intermediate_dim]
    889
            }
    890
    10
        }
    891
    892
        /// Return an intermediate buffer to the pool
    893
    10
        pub fn return_intermediate(&self, mut buffer: Vec<f32>) {
    894
    10
            self.returns
    895
    10
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    896
    10
            buffer.fill(0.0);
    897
    898
    10
            let mut buffers = self.intermediate_buffers.lock().expect("mutex poisoned");
    899
    10
            if buffers.len() < self.pool_size {
    900
    10
                buffers.push(buffer);
    901
    10
            }0
    902
    10
        }
    903
    904
        /// Borrow an attention score buffer from the pool
    905
    10
        pub fn borrow_attention(&self) -> Vec<f32> {
    906
    10
            self.borrows
    907
    10
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    908
    909
    10
            let mut buffers = self.attention_buffers.lock().expect("mutex poisoned");
    910
    10
            if let Some(buffer) = buffers.pop() {
    911
    10
                buffer
    912
            } else {
    913
    0
                if self.warmed_up.load(std::sync::atomic::Ordering::Acquire) {
    914
    0
                    self.post_warmup_allocs
    915
    0
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    916
    0
                }
    917
    0
                vec![0.0f32; self.num_heads * self.max_seq_len]
    918
            }
    919
    10
        }
    920
    921
        /// Return an attention score buffer to the pool
    922
    10
        pub fn return_attention(&self, mut buffer: Vec<f32>) {
    923
    10
            self.returns
    924
    10
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    925
    10
            buffer.fill(0.0);
    926
    927
    10
            let mut buffers = self.attention_buffers.lock().expect("mutex poisoned");
    928
    10
            if buffers.len() < self.pool_size {
    929
    10
                buffers.push(buffer);
    930
    10
            }0
    931
    10
        }
    932
    933
        /// Check if pool has achieved zero-allocation after warmup
    934
    1
        pub fn is_zero_alloc(&self) -> bool {
    935
    1
            self.warmed_up.load(std::sync::atomic::Ordering::Acquire)
    936
    1
                && self
    937
    1
                    .post_warmup_allocs
    938
    1
                    .load(std::sync::atomic::Ordering::Relaxed)
    939
    1
                    == 0
    940
    1
        }
    941
    942
        /// Get pool statistics
    943
    6
        pub fn stats(&self) -> GpuBufferPoolStats {
    944
    6
            GpuBufferPoolStats {
    945
    6
                borrows: self.borrows.load(std::sync::atomic::Ordering::Relaxed),
    946
    6
                returns: self.returns.load(std::sync::atomic::Ordering::Relaxed),
    947
    6
                post_warmup_allocs: self
    948
    6
                    .post_warmup_allocs
    949
    6
                    .load(std::sync::atomic::Ordering::Relaxed),
    950
    6
                warmed_up: self.warmed_up.load(std::sync::atomic::Ordering::Acquire),
    951
    6
                hidden_available: self.hidden_buffers.lock().expect("mutex poisoned").len(),
    952
    6
                intermediate_available: self
    953
    6
                    .intermediate_buffers
    954
    6
                    .lock()
    955
    6
                    .expect("mutex poisoned")
    956
    6
                    .len(),
    957
    6
                attention_available: self.attention_buffers.lock().expect("mutex poisoned").len(),
    958
    6
            }
    959
    6
        }
    960
    961
        /// Calculate total memory usage of the buffer pool
    962
    1
        pub fn memory_usage_bytes(&self) -> usize {
    963
    1
            let hidden_bytes = self.pool_size * self.hidden_dim * 4;
    964
    1
            let intermediate_bytes = self.pool_size * self.intermediate_dim * 4;
    965
    1
            let attention_bytes = self.pool_size * self.num_heads * self.max_seq_len * 4;
    966
    1
            hidden_bytes + intermediate_bytes + attention_bytes
    967
    1
        }
    968
    }
    969
    970
    /// Statistics for GpuBufferPool
    971
    #[cfg(feature = "gpu")]
    972
    #[derive(Debug, Clone)]
    973
    pub struct GpuBufferPoolStats {
    974
        /// Total borrows
    975
        pub borrows: u64,
    976
        /// Total returns
    977
        pub returns: u64,
    978
        /// Allocations after warmup (should be 0)
    979
        pub post_warmup_allocs: u64,
    980
        /// Whether warmup is complete
    981
        pub warmed_up: bool,
    982
        /// Available hidden buffers
    983
        pub hidden_available: usize,
    984
        /// Available intermediate buffers
    985
        pub intermediate_available: usize,
    986
        /// Available attention buffers
    987
        pub attention_available: usize,
    988
    }
    989
    990
    /// Async Command Queue for GPU pipelining (PARITY-032, IMP-310)
    991
    ///
    992
    /// Implements double-buffering to hide GPU latency by overlapping
    993
    /// computation and data transfer. While one batch is being processed
    994
    /// on GPU, the next batch is being prepared on CPU.
    995
    ///
    996
    /// # Key Properties
    997
    /// - Double-buffering: 2 command slots for overlap
    998
    /// - Async submission: Non-blocking command enqueue
    999
    /// - Pipeline stages: Prepare → Submit → Execute → Complete
    1000
    ///
    1001
    /// # GPU Utilization Target
    1002
    /// - Without pipelining: ~50% (waiting for results)
    1003
    /// - With pipelining: >85% (overlapped execution)
    1004
    #[cfg(feature = "gpu")]
    1005
    pub struct AsyncCommandQueue {
    1006
        /// Command slots for double-buffering (2 slots)
    1007
        slots: [std::sync::Mutex<CommandSlot>; 2],
    1008
        /// Current slot index for submission
    1009
        current_slot: std::sync::atomic::AtomicUsize,
    1010
        /// Statistics: commands submitted
    1011
        pub commands_submitted: std::sync::atomic::AtomicU64,
    1012
        /// Statistics: commands completed
    1013
        pub commands_completed: std::sync::atomic::AtomicU64,
    1014
        /// Statistics: pipeline stalls (had to wait for previous)
    1015
        pub pipeline_stalls: std::sync::atomic::AtomicU64,
    1016
    }
    1017
    1018
    /// State of a command slot in the async queue
    1019
    #[cfg(feature = "gpu")]
    1020
    #[derive(Debug, Clone)]
    1021
    pub enum CommandSlotState {
    1022
        /// Slot is empty and ready for new command
    1023
        Empty,
    1024
        /// Command is being prepared (CPU side)
    1025
        Preparing,
    1026
        /// Command has been submitted to GPU
    1027
        Submitted,
    1028
        /// Command execution is complete
    1029
        Complete,
    1030
    }
    1031
    1032
    /// A command slot for async execution
    1033
    #[cfg(feature = "gpu")]
    1034
    pub struct CommandSlot {
    1035
        /// Current state of this slot
    1036
        state: CommandSlotState,
    1037
        /// Input data for the command
    1038
        input: Option<Vec<f32>>,
    1039
        /// Output data from the command
    1040
        output: Option<Vec<f32>>,
    1041
        /// Timestamp when command was submitted
    1042
        submit_time: Option<std::time::Instant>,
    1043
    }
    1044
    1045
    #[cfg(feature = "gpu")]
    1046
    impl Default for CommandSlot {
    1047
    8
        fn default() -> Self {
    1048
    8
            Self {
    1049
    8
                state: CommandSlotState::Empty,
    1050
    8
                input: None,
    1051
    8
                output: None,
    1052
    8
                submit_time: None,
    1053
    8
            }
    1054
    8
        }
    1055
    }
    1056
    1057
    #[cfg(feature = "gpu")]
    1058
    impl AsyncCommandQueue {
    1059
        /// Create new async command queue with double-buffering
    1060
    4
        pub fn new() -> Self {
    1061
    4
            Self {
    1062
    4
                slots: [
    1063
    4
                    std::sync::Mutex::new(CommandSlot::default()),
    1064
    4
                    std::sync::Mutex::new(CommandSlot::default()),
    1065
    4
                ],
    1066
    4
                current_slot: std::sync::atomic::AtomicUsize::new(0),
    1067
    4
                commands_submitted: std::sync::atomic::AtomicU64::new(0),
    1068
    4
                commands_completed: std::sync::atomic::AtomicU64::new(0),
    1069
    4
                pipeline_stalls: std::sync::atomic::AtomicU64::new(0),
    1070
    4
            }
    1071
    4
        }
    1072
    1073
        /// Submit a command for async execution
    1074
        ///
    1075
        /// Returns the slot index where the command was placed.
    1076
        /// If both slots are busy, this will block until one is available
    1077
        /// (counted as a pipeline stall).
    1078
    24
        pub fn submit(&self, input: Vec<f32>) -> usize {
    1079
    24
            let slot_idx = self
    1080
    24
                .current_slot
    1081
    24
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
    1082
    24
                % 2;
    1083
    1084
    24
            let mut slot = self.slots[slot_idx].lock().expect("mutex poisoned");
    1085
    1086
            // Check if we need to wait for previous command
    1087
    24
            if matches!(
    1088
    24
                slot.state,
    1089
                CommandSlotState::Submitted | CommandSlotState::Preparing
    1090
    0
            ) {
    1091
    0
                self.pipeline_stalls
    1092
    0
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1093
    0
                // In real implementation, would wait for GPU completion
    1094
    0
                // For now, mark as complete to allow reuse
    1095
    0
                slot.state = CommandSlotState::Complete;
    1096
    24
            }
    1097
    1098
            // Prepare new command
    1099
    24
            slot.state = CommandSlotState::Preparing;
    1100
    24
            slot.input = Some(input);
    1101
    24
            slot.output = None;
    1102
    24
            slot.submit_time = Some(std::time::Instant::now());
    1103
    1104
            // Mark as submitted
    1105
    24
            slot.state = CommandSlotState::Submitted;
    1106
    24
            self.commands_submitted
    1107
    24
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1108
    1109
    24
            slot_idx
    1110
    24
        }
    1111
    1112
        /// Mark a command as complete with output
    1113
    22
        pub fn complete(&self, slot_idx: usize, output: Vec<f32>) {
    1114
    22
            let mut slot = self.slots[slot_idx].lock().expect("mutex poisoned");
    1115
    22
            slot.state = CommandSlotState::Complete;
    1116
    22
            slot.output = Some(output);
    1117
    22
            self.commands_completed
    1118
    22
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1119
    22
        }
    1120
    1121
        /// Get output from a completed command
    1122
        ///
    1123
        /// Returns None if command is not complete yet.
    1124
    21
        pub fn get_output(&self, slot_idx: usize) -> Option<Vec<f32>> {
    1125
    21
            let mut slot = self.slots[slot_idx].lock().expect("mutex poisoned");
    1126
    21
            if matches!0(slot.state, CommandSlotState::Complete) {
    1127
    21
                slot.state = CommandSlotState::Empty;
    1128
    21
                slot.output.take()
    1129
            } else {
    1130
    0
                None
    1131
            }
    1132
    21
        }
    1133
    1134
        /// Get queue statistics
    1135
    5
        pub fn stats(&self) -> AsyncQueueStats {
    1136
    5
            let submitted = self
    1137
    5
                .commands_submitted
    1138
    5
                .load(std::sync::atomic::Ordering::Relaxed);
    1139
    5
            let completed = self
    1140
    5
                .commands_completed
    1141
    5
                .load(std::sync::atomic::Ordering::Relaxed);
    1142
    5
            let stalls = self
    1143
    5
                .pipeline_stalls
    1144
    5
                .load(std::sync::atomic::Ordering::Relaxed);
    1145
    1146
            // GPU utilization estimate: (1 - stalls/submitted) * 100
    1147
    5
            let utilization = if submitted > 0 {
    1148
    4
                (1.0 - stalls as f64 / submitted as f64) * 100.0
    1149
            } else {
    1150
    1
                0.0
    1151
            };
    1152
    1153
    5
            AsyncQueueStats {
    1154
    5
                commands_submitted: submitted,
    1155
    5
                commands_completed: completed,
    1156
    5
                pipeline_stalls: stalls,
    1157
    5
                in_flight: submitted.saturating_sub(completed),
    1158
    5
                gpu_utilization_percent: utilization,
    1159
    5
            }
    1160
    5
        }
    1161
    1162
        /// Calculate pipeline efficiency
    1163
        ///
    1164
        /// Efficiency = commands without stall / total commands
    1165
    1
        pub fn pipeline_efficiency(&self) -> f64 {
    1166
    1
            let submitted = self
    1167
    1
                .commands_submitted
    1168
    1
                .load(std::sync::atomic::Ordering::Relaxed);
    1169
    1
            let stalls = self
    1170
    1
                .pipeline_stalls
    1171
    1
                .load(std::sync::atomic::Ordering::Relaxed);
    1172
    1
            if submitted == 0 {
    1173
    0
                return 1.0;
    1174
    1
            }
    1175
    1
            (submitted - stalls) as f64 / submitted as f64
    1176
    1
        }
    1177
    }
    1178
    1179
    #[cfg(feature = "gpu")]
    1180
    impl Default for AsyncCommandQueue {
    1181
    0
        fn default() -> Self {
    1182
    0
            Self::new()
    1183
    0
        }
    1184
    }
    1185
    1186
    /// Statistics for AsyncCommandQueue
    1187
    #[cfg(feature = "gpu")]
    1188
    #[derive(Debug, Clone)]
    1189
    pub struct AsyncQueueStats {
    1190
        /// Total commands submitted
    1191
        pub commands_submitted: u64,
    1192
        /// Total commands completed
    1193
        pub commands_completed: u64,
    1194
        /// Pipeline stalls (had to wait)
    1195
        pub pipeline_stalls: u64,
    1196
        /// Commands currently in flight
    1197
        pub in_flight: u64,
    1198
        /// Estimated GPU utilization percentage
    1199
        pub gpu_utilization_percent: f64,
    1200
    }
    1201
    1202
    /// Prefix Cache for common prompts (PARITY-033, IMP-319)
    1203
    ///
    1204
    /// Caches the KV cache state for common prompt prefixes, enabling
    1205
    /// instant response (0ms TTFT) for repeated prompts.
    1206
    ///
    1207
    /// # Key Properties
    1208
    /// - Hash-based prefix lookup (FNV-1a)
    1209
    /// - LRU eviction for memory management
    1210
    /// - Thread-safe access
    1211
    ///
    1212
    /// # Use Cases
    1213
    /// - System prompts (cached once, reused for all requests)
    1214
    /// - Common few-shot examples
    1215
    /// - Chat history prefixes
    1216
    #[cfg(feature = "gpu")]
    1217
    pub struct PrefixCache {
    1218
        /// Cached prefix entries (hash → entry)
    1219
        entries: std::sync::Mutex<std::collections::HashMap<u64, PrefixCacheEntry>>,
    1220
        /// Maximum number of cached prefixes
    1221
        max_entries: usize,
    1222
        /// Statistics: cache hits
    1223
        pub hits: std::sync::atomic::AtomicU64,
    1224
        /// Statistics: cache misses
    1225
        pub misses: std::sync::atomic::AtomicU64,
    1226
        /// Statistics: evictions
    1227
        pub evictions: std::sync::atomic::AtomicU64,
    1228
    }
    1229
    1230
    /// A cached prefix entry
    1231
    #[cfg(feature = "gpu")]
    1232
    pub struct PrefixCacheEntry {
    1233
        /// The original prompt tokens
    1234
        pub tokens: Vec<u32>,
    1235
        /// Cached K state for each layer [num_layers, seq_len, hidden_dim]
    1236
        pub k_cache: Vec<Vec<f32>>,
    1237
        /// Cached V state for each layer [num_layers, seq_len, hidden_dim]
    1238
        pub v_cache: Vec<Vec<f32>>,
    1239
        /// Timestamp for LRU eviction
    1240
        pub last_access: std::time::Instant,
    1241
        /// Number of times this prefix was hit
    1242
        pub hit_count: u64,
    1243
    }
    1244
    1245
    #[cfg(feature = "gpu")]
    1246
    impl PrefixCache {
    1247
        /// Create new prefix cache with specified capacity
    1248
    4
        pub fn new(max_entries: usize) -> Self {
    1249
    4
            Self {
    1250
    4
                entries: std::sync::Mutex::new(std::collections::HashMap::with_capacity(max_entries)),
    1251
    4
                max_entries,
    1252
    4
                hits: std::sync::atomic::AtomicU64::new(0),
    1253
    4
                misses: std::sync::atomic::AtomicU64::new(0),
    1254
    4
                evictions: std::sync::atomic::AtomicU64::new(0),
    1255
    4
            }
    1256
    4
        }
    1257
    1258
        /// Hash tokens to create cache key (FNV-1a)
    1259
    11
        fn hash_tokens(tokens: &[u32]) -> u64 {
    1260
            const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    1261
            const FNV_PRIME: u64 = 0x0100_0000_01b3;
    1262
    1263
    11
            let mut hash = FNV_OFFSET;
    1264
    287
            for &token276 in tokens {
    1265
    276
                hash ^= token as u64;
    1266
    276
                hash = hash.wrapping_mul(FNV_PRIME);
    1267
    276
            }
    1268
    11
            hash
    1269
    11
        }
    1270
    1271
        /// Look up a prefix in the cache
    1272
        ///
    1273
        /// Returns the cached KV state if found, None otherwise.
    1274
        #[allow(clippy::type_complexity)]
    1275
    5
        pub fn lookup(&self, tokens: &[u32]) -> Option<(Vec<Vec<f32>>, Vec<Vec<f32>>)> {
    1276
    5
            let hash = Self::hash_tokens(tokens);
    1277
    1278
    5
            let mut entries = self.entries.lock().expect("mutex poisoned");
    1279
    5
            if let Some(entry3) = entries.get_mut(&hash) {
    1280
                // Verify tokens match (hash collision check)
    1281
    3
                if entry.tokens == tokens {
    1282
    3
                    self.hits.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1283
    3
                    entry.last_access = std::time::Instant::now();
    1284
    3
                    entry.hit_count += 1;
    1285
    3
                    return Some((entry.k_cache.clone(), entry.v_cache.clone()));
    1286
    0
                }
    1287
    2
            }
    1288
    1289
    2
            self.misses
    1290
    2
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1291
    2
            None
    1292
    5
        }
    1293
    1294
        /// Insert a new prefix into the cache
    1295
        ///
    1296
        /// Evicts LRU entry if cache is full.
    1297
    6
        pub fn insert(&self, tokens: Vec<u32>, k_cache: Vec<Vec<f32>>, v_cache: Vec<Vec<f32>>) {
    1298
    6
            let hash = Self::hash_tokens(&tokens);
    1299
    1300
    6
            let mut entries = self.entries.lock().expect("mutex poisoned");
    1301
    1302
            // Evict LRU if at capacity
    1303
    6
            if entries.len() >= self.max_entries {
    1304
                // Find oldest entry
    1305
    1
                if let Some((&oldest_hash, _)) = entries.iter().min_by_key(|(_, e)| e.last_access) {
    1306
    1
                    entries.remove(&oldest_hash);
    1307
    1
                    self.evictions
    1308
    1
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1309
    1
                }0
    1310
    5
            }
    1311
    1312
    6
            entries.insert(
    1313
    6
                hash,
    1314
    6
                PrefixCacheEntry {
    1315
    6
                    tokens,
    1316
    6
                    k_cache,
    1317
    6
                    v_cache,
    1318
    6
                    last_access: std::time::Instant::now(),
    1319
    6
                    hit_count: 0,
    1320
    6
                },
    1321
    6
            );
    1322
    6
        }
    1323
    1324
        /// Check if a prefix is cached
    1325
    0
        pub fn contains(&self, tokens: &[u32]) -> bool {
    1326
    0
            let hash = Self::hash_tokens(tokens);
    1327
    0
            let entries = self.entries.lock().expect("mutex poisoned");
    1328
    0
            entries.contains_key(&hash)
    1329
    0
        }
    1330
    1331
        /// Get cache statistics
    1332
    6
        pub fn stats(&self) -> PrefixCacheStats {
    1333
    6
            let hits = self.hits.load(std::sync::atomic::Ordering::Relaxed);
    1334
    6
            let misses = self.misses.load(std::sync::atomic::Ordering::Relaxed);
    1335
    6
            let total = hits + misses;
    1336
    1337
            PrefixCacheStats {
    1338
    6
                hits,
    1339
    6
                misses,
    1340
    6
                evictions: self.evictions.load(std::sync::atomic::Ordering::Relaxed),
    1341
    6
                entries: self.entries.lock().expect("mutex poisoned").len(),
    1342
    6
                hit_rate: if total > 0 {
    1343
    3
                    hits as f64 / total as f64
    1344
                } else {
    1345
    3
                    0.0
    1346
                },
    1347
            }
    1348
    6
        }
    1349
    1350
        /// Clear all cached entries
    1351
    0
        pub fn clear(&self) {
    1352
    0
            let mut entries = self.entries.lock().expect("mutex poisoned");
    1353
    0
            entries.clear();
    1354
    0
        }
    1355
    1356
        /// Estimate memory usage of cached prefixes
    1357
    1
        pub fn memory_usage_bytes(&self) -> usize {
    1358
    1
            let entries = self.entries.lock().expect("mutex poisoned");
    1359
    1
            entries
    1360
    1
                .values()
    1361
    1
                .map(|e| {
    1362
    32
                    let k_bytes1: usize1 = e.k_cache.iter()1.map1(|v| v.len() * 4).sum1();
    1363
    32
                    let v_bytes1: usize1 = e.v_cache.iter()1.map1(|v| v.len() * 4).sum1();
    1364
    1
                    let token_bytes = e.tokens.len() * 4;
    1365
    1
                    k_bytes + v_bytes + token_bytes
    1366
    1
                })
    1367
    1
                .sum()
    1368
    1
        }
    1369
    }
    1370
    1371
    #[cfg(feature = "gpu")]
    1372
    impl Default for PrefixCache {
    1373
    0
        fn default() -> Self {
    1374
    0
            Self::new(16) // Default: cache 16 prefixes
    1375
    0
        }
    1376
    }
    1377
    1378
    /// Statistics for PrefixCache
    1379
    #[cfg(feature = "gpu")]
    1380
    #[derive(Debug, Clone)]
    1381
    pub struct PrefixCacheStats {
    1382
        /// Cache hits
    1383
        pub hits: u64,
    1384
        /// Cache misses
    1385
        pub misses: u64,
    1386
        /// Evictions due to capacity
    1387
        pub evictions: u64,
    1388
        /// Current number of cached entries
    1389
        pub entries: usize,
    1390
        /// Hit rate (0.0 - 1.0)
    1391
        pub hit_rate: f64,
    1392
    }
    1393
    1394
    // =============================================================================
    1395
    // PARITY-034: Multi-Request Scheduler with Scheduling Policies (IMP-317)
    1396
    // =============================================================================
    1397
    //
    1398
    // Extends PARITY-028's ContinuousBatchScheduler with:
    1399
    // - Multiple scheduling policies (FCFS, SJF, Round-Robin)
    1400
    // - Request queuing with priorities
    1401
    // - TTFT (Time to First Token) tracking
    1402
    // - Throughput scaling verification
    1403
    //
    1404
    // Architecture:
    1405
    // - Incoming requests are queued with their KV cache states
    1406
    // - Scheduler batches decode steps from multiple requests
    1407
    // - GPU GEMM efficiency: batch_size > 1 enables GPU acceleration
    1408
    // - Preemption: Long-running requests can be paused for new arrivals
    1409
    // =============================================================================
    1410
    1411
    /// Request state in the multi-request scheduler
    1412
    #[cfg(feature = "gpu")]
    1413
    #[derive(Debug, Clone, PartialEq, Eq)]
    1414
    pub enum MultiRequestState {
    1415
        /// Waiting for prefill
    1416
        Pending,
    1417
        /// Prefill in progress
    1418
        Prefilling,
    1419
        /// Decode in progress
    1420
        Decoding,
    1421
        /// Request completed
    1422
        Completed,
    1423
        /// Request preempted (paused)
    1424
        Preempted,
    1425
    }
    1426
    1427
    /// A single inference request in the multi-request scheduler
    1428
    #[cfg(feature = "gpu")]
    1429
    #[derive(Clone)]
    1430
    pub struct MultiSchedulerRequest {
    1431
        /// Unique request ID
    1432
        pub id: u64,
    1433
        /// Input tokens
    1434
        pub tokens: Vec<u32>,
    1435
        /// Generated tokens so far
    1436
        pub generated: Vec<u32>,
    1437
        /// Maximum tokens to generate
    1438
        pub max_tokens: usize,
    1439
        /// Current state
    1440
        pub state: MultiRequestState,
    1441
        /// KV cache position (how many tokens processed)
    1442
        pub kv_position: usize,
    1443
        /// Arrival time for FCFS scheduling
    1444
        pub arrival_time: std::time::Instant,
    1445
        /// Time first token generated (for TTFT metric)
    1446
        pub first_token_time: Option<std::time::Instant>,
    1447
    }
    1448
    1449
    #[cfg(feature = "gpu")]
    1450
    impl MultiSchedulerRequest {
    1451
        /// Create new request
    1452
    22
        pub fn new(id: u64, tokens: Vec<u32>, max_tokens: usize) -> Self {
    1453
    22
            Self {
    1454
    22
                id,
    1455
    22
                tokens,
    1456
    22
                generated: Vec::with_capacity(max_tokens),
    1457
    22
                max_tokens,
    1458
    22
                state: MultiRequestState::Pending,
    1459
    22
                kv_position: 0,
    1460
    22
                arrival_time: std::time::Instant::now(),
    1461
    22
                first_token_time: None,
    1462
    22
            }
    1463
    22
        }
    1464
    1465
        /// Check if request is complete
    1466
    503
        pub fn is_complete(&self) -> bool {
    1467
    503
            self.state == MultiRequestState::Completed || self.generated.len() >= self.max_tokens
    1468
    503
        }
    1469
    1470
        /// Time to first token (None if not yet generated)
    1471
    0
        pub fn ttft_ms(&self) -> Option<f64> {
    1472
    0
            self.first_token_time
    1473
    0
                .map(|t| t.duration_since(self.arrival_time).as_secs_f64() * 1000.0)
    1474
    0
        }
    1475
    }
    1476
    1477
    /// Scheduling policy for the batch scheduler
    1478
    #[cfg(feature = "gpu")]
    1479
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    1480
    pub enum SchedulingPolicy {
    1481
        /// First-come, first-served
    1482
        Fcfs,
    1483
        /// Shortest job first (by remaining tokens)
    1484
        Sjf,
    1485
        /// Round-robin with time slices
    1486
        RoundRobin,
    1487
    }
    1488
    1489
    /// Multi-request scheduler with scheduling policies (PARITY-034)
    1490
    #[cfg(feature = "gpu")]
    1491
    pub struct MultiRequestScheduler {
    1492
        /// Pending requests queue
    1493
        pending: std::sync::Mutex<std::collections::VecDeque<MultiSchedulerRequest>>,
    1494
        /// Active requests being processed
    1495
        active: std::sync::Mutex<Vec<MultiSchedulerRequest>>,
    1496
        /// Completed requests
    1497
        completed: std::sync::Mutex<Vec<MultiSchedulerRequest>>,
    1498
        /// Maximum batch size
    1499
        max_batch_size: usize,
    1500
        /// Maximum concurrent requests
    1501
        max_concurrent: usize,
    1502
        /// Scheduling policy
    1503
        policy: SchedulingPolicy,
    1504
        /// Request ID counter
    1505
        next_id: std::sync::atomic::AtomicU64,
    1506
        /// Requests submitted
    1507
        pub requests_submitted: std::sync::atomic::AtomicU64,
    1508
        /// Requests completed
    1509
        pub requests_completed: std::sync::atomic::AtomicU64,
    1510
        /// Total tokens generated
    1511
        pub tokens_generated: std::sync::atomic::AtomicU64,
    1512
        /// Batch iterations performed
    1513
        pub batch_iterations: std::sync::atomic::AtomicU64,
    1514
    }
    1515
    1516
    #[cfg(feature = "gpu")]
    1517
    impl MultiRequestScheduler {
    1518
        /// Create new scheduler with given parameters
    1519
    7
        pub fn new(max_batch_size: usize, max_concurrent: usize, policy: SchedulingPolicy) -> Self {
    1520
    7
            Self {
    1521
    7
                pending: std::sync::Mutex::new(std::collections::VecDeque::new()),
    1522
    7
                active: std::sync::Mutex::new(Vec::with_capacity(max_concurrent)),
    1523
    7
                completed: std::sync::Mutex::new(Vec::new()),
    1524
    7
                max_batch_size,
    1525
    7
                max_concurrent,
    1526
    7
                policy,
    1527
    7
                next_id: std::sync::atomic::AtomicU64::new(0),
    1528
    7
                requests_submitted: std::sync::atomic::AtomicU64::new(0),
    1529
    7
                requests_completed: std::sync::atomic::AtomicU64::new(0),
    1530
    7
                tokens_generated: std::sync::atomic::AtomicU64::new(0),
    1531
    7
                batch_iterations: std::sync::atomic::AtomicU64::new(0),
    1532
    7
            }
    1533
    7
        }
    1534
    1535
        /// Submit a new request
    1536
    22
        pub fn submit(&self, tokens: Vec<u32>, max_tokens: usize) -> u64 {
    1537
    22
            let id = self
    1538
    22
                .next_id
    1539
    22
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1540
    22
            let request = MultiSchedulerRequest::new(id, tokens, max_tokens);
    1541
    1542
    22
            let mut pending = self.pending.lock().expect("mutex poisoned");
    1543
    22
            pending.push_back(request);
    1544
    22
            self.requests_submitted
    1545
    22
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1546
    1547
    22
            id
    1548
    22
        }
    1549
    1550
        /// Get batch of requests ready for decode step
    1551
        ///
    1552
        /// Returns request IDs and their current positions
    1553
    107
        pub fn get_decode_batch(&self) -> Vec<(u64, usize)> {
    1554
    107
            let mut active = self.active.lock().expect("mutex poisoned");
    1555
    107
            let mut pending = self.pending.lock().expect("mutex poisoned");
    1556
    1557
            // Promote pending requests to active (up to max_concurrent)
    1558
    129
            while active.len() < self.max_concurrent && !pending.is_empty() {
    1559
    22
                if let Some(mut req) = pending.pop_front() {
    1560
    22
                    req.state = MultiRequestState::Decoding;
    1561
    22
                    active.push(req);
    1562
    22
                }0
    1563
            }
    1564
    1565
            // Sort by policy
    1566
    107
            match self.policy {
    1567
    103
                SchedulingPolicy::Fcfs => {
    1568
    103
                    // Already in arrival order
    1569
    103
                },
    1570
                SchedulingPolicy::Sjf => {
    1571
    10
                    active2.sort_by_key2(|r| r.max_tokens - r.generated.len());
    1572
                },
    1573
                SchedulingPolicy::RoundRobin => {
    1574
                    // Rotate - move first to end
    1575
    2
                    if active.len() > 1 {
    1576
    2
                        let first = active.remove(0);
    1577
    2
                        active.push(first);
    1578
    2
                    }0
    1579
                },
    1580
            }
    1581
    1582
            // Return batch of decoding requests
    1583
    107
            active
    1584
    107
                .iter()
    1585
    514
                .filter107(|r| r.state == MultiRequestState::Decoding)
    1586
    107
                .take(self.max_batch_size)
    1587
    514
                .map107(|r| (r.id, r.kv_position))
    1588
    107
                .collect()
    1589
    107
        }
    1590
    1591
        /// Record generated token for a request
    1592
    503
        pub fn record_token(&self, request_id: u64, token: u32) {
    1593
    503
            let mut active = self.active.lock().expect("mutex poisoned");
    1594
    1595
    1.95k
            if let Some(req503) = active.iter_mut()503.find503(|r| r.id == request_id) {
    1596
                // Record TTFT for first token
    1597
    503
                if req.first_token_time.is_none() {
    1598
    11
                    req.first_token_time = Some(std::time::Instant::now());
    1599
    492
                }
    1600
    1601
    503
                req.generated.push(token);
    1602
    503
                req.kv_position += 1;
    1603
    503
                self.tokens_generated
    1604
    503
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1605
    1606
                // Check if complete
    1607
    503
                if req.is_complete() {
    1608
    11
                    req.state = MultiRequestState::Completed;
    1609
    492
                }
    1610
    0
            }
    1611
    503
        }
    1612
    1613
        /// Move completed requests from active to completed
    1614
    101
        pub fn collect_completed(&self) -> Vec<MultiSchedulerRequest> {
    1615
    101
            let mut active = self.active.lock().expect("mutex poisoned");
    1616
    101
            let mut completed = self.completed.lock().expect("mutex poisoned");
    1617
    1618
    101
            let (done, still_active): (Vec<_>, Vec<_>) = active
    1619
    101
                .drain(..)
    1620
    601
                .partition101(|r| r.state == MultiRequestState::Completed);
    1621
    1622
    101
            *active = still_active;
    1623
    1624
    112
            for _req11 in &done {
    1625
    11
                self.requests_completed
    1626
    11
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1627
    11
            }
    1628
    1629
    101
            completed.extend(done.iter().cloned());
    1630
    101
            done
    1631
    101
        }
    1632
    1633
        /// Run one batch iteration (for simulation)
    1634
    103
        pub fn step(&self) {
    1635
    103
            self.batch_iterations
    1636
    103
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    1637
    103
        }
    1638
    1639
        /// Get scheduler statistics
    1640
    107
        pub fn stats(&self) -> MultiRequestStats {
    1641
    107
            let submitted = self
    1642
    107
                .requests_submitted
    1643
    107
                .load(std::sync::atomic::Ordering::Relaxed);
    1644
    107
            let completed = self
    1645
    107
                .requests_completed
    1646
    107
                .load(std::sync::atomic::Ordering::Relaxed);
    1647
    107
            let tokens = self
    1648
    107
                .tokens_generated
    1649
    107
                .load(std::sync::atomic::Ordering::Relaxed);
    1650
    107
            let iterations = self
    1651
    107
                .batch_iterations
    1652
    107
                .load(std::sync::atomic::Ordering::Relaxed);
    1653
    1654
    107
            let pending = self.pending.lock().expect("mutex poisoned").len();
    1655
    107
            let active = self.active.lock().expect("mutex poisoned").len();
    1656
    1657
            MultiRequestStats {
    1658
    107
                requests_submitted: submitted,
    1659
    107
                requests_completed: completed,
    1660
    107
                tokens_generated: tokens,
    1661
    107
                batch_iterations: iterations,
    1662
    107
                pending_requests: pending,
    1663
    107
                active_requests: active,
    1664
    107
                avg_batch_size: if iterations > 0 {
    1665
    103
                    tokens as f64 / iterations as f64
    1666
                } else {
    1667
    4
                    0.0
    1668
                },
    1669
            }
    1670
    107
        }
    1671
    }
    1672
    1673
    /// Statistics for multi-request scheduler (PARITY-034)
    1674
    #[cfg(feature = "gpu")]
    1675
    pub struct MultiRequestStats {
    1676
        /// Total requests submitted
    1677
        pub requests_submitted: u64,
    1678
        /// Total requests completed
    1679
        pub requests_completed: u64,
    1680
        /// Total tokens generated
    1681
        pub tokens_generated: u64,
    1682
        /// Batch iterations performed
    1683
        pub batch_iterations: u64,
    1684
        /// Current pending requests
    1685
        pub pending_requests: usize,
    1686
        /// Current active requests
    1687
        pub active_requests: usize,
    1688
        /// Average batch size
    1689
        pub avg_batch_size: f64,
    1690
    }
    1691
    1692
    // =============================================================================
    1693
    // PARITY-035: Chunked Prefill for Long Contexts (IMP-320)
    1694
    // =============================================================================
    1695
    //
    1696
    // Enables streaming prompt processing by breaking long prefills into chunks.
    1697
    // Key optimization for TTFT (Time to First Token) with long contexts.
    1698
    //
    1699
    // Architecture:
    1700
    // - Prompt is split into chunks (default 512 tokens)
    1701
    // - Each chunk processes incrementally, updating KV cache
    1702
    // - First token can be generated after first chunk completes
    1703
    // - Total prefill time is spread across chunks
    1704
    // =============================================================================
    1705
    1706
    /// Configuration for chunked prefill
    1707
    #[cfg(feature = "gpu")]
    1708
    #[derive(Debug, Clone)]
    1709
    pub struct ChunkedPrefillConfig {
    1710
        /// Chunk size in tokens (default: 512)
    1711
        pub chunk_size: usize,
    1712
        /// Maximum context length (default: 8192)
    1713
        pub max_context: usize,
    1714
        /// Whether to yield after each chunk for streaming
    1715
        pub stream_chunks: bool,
    1716
    }
    1717
    1718
    #[cfg(feature = "gpu")]
    1719
    impl Default for ChunkedPrefillConfig {
    1720
    5
        fn default() -> Self {
    1721
    5
            Self {
    1722
    5
                chunk_size: 512,
    1723
    5
                max_context: 8192,
    1724
    5
                stream_chunks: true,
    1725
    5
            }
    1726
    5
        }
    1727
    }
    1728
    1729
    #[cfg(feature = "gpu")]
    1730
    impl ChunkedPrefillConfig {
    1731
        /// Create config with custom chunk size
    1732
    5
        pub fn with_chunk_size(chunk_size: usize) -> Self {
    1733
    5
            Self {
    1734
    5
                chunk_size,
    1735
    5
                ..Default::default()
    1736
    5
            }
    1737
    5
        }
    1738
    }
    1739
    1740
    /// Progress report for a single chunk
    1741
    #[cfg(feature = "gpu")]
    1742
    #[derive(Debug, Clone)]
    1743
    pub struct ChunkProgress {
    1744
        /// Chunk index (0-based)
    1745
        pub chunk_idx: usize,
    1746
        /// Total chunks
    1747
        pub total_chunks: usize,
    1748
        /// Tokens processed so far
    1749
        pub tokens_processed: usize,
    1750
        /// Total tokens to process
    1751
        pub total_tokens: usize,
    1752
        /// Time for this chunk (ms)
    1753
        pub chunk_time_ms: f64,
    1754
        /// Cumulative time so far (ms)
    1755
        pub cumulative_time_ms: f64,
    1756
    }
    1757
    1758
    /// Chunked prefill processor for long context handling
    1759
    #[cfg(feature = "gpu")]
    1760
    pub struct ChunkedPrefill {
    1761
        /// Configuration
    1762
        config: ChunkedPrefillConfig,
    1763
        /// Chunks created from prompt
    1764
        chunks: Vec<Vec<u32>>,
    1765
        /// Current chunk being processed
    1766
        current_chunk: usize,
    1767
        /// Tokens processed so far
    1768
        tokens_processed: usize,
    1769
        /// Start time for timing
    1770
        start_time: Option<std::time::Instant>,
    1771
        /// Timing for each chunk
    1772
        chunk_times_ms: Vec<f64>,
    1773
    }
    1774
    1775
    #[cfg(feature = "gpu")]
    1776
    impl ChunkedPrefill {
    1777
        /// Create new chunked prefill from prompt tokens
    1778
    5
        pub fn new(prompt_tokens: &[u32], config: ChunkedPrefillConfig) -> Self {
    1779
    5
            let chunks: Vec<Vec<u32>> = prompt_tokens
    1780
    5
                .chunks(config.chunk_size)
    1781
    5
                .map(<[u32]>::to_vec)
    1782
    5
                .collect();
    1783
    1784
    5
            Self {
    1785
    5
                config,
    1786
    5
                chunks,
    1787
    5
                current_chunk: 0,
    1788
    5
                tokens_processed: 0,
    1789
    5
                start_time: None,
    1790
    5
                chunk_times_ms: Vec::new(),
    1791
    5
            }
    1792
    5
        }
    1793
    1794
        /// Get total number of chunks
    1795
    3
        pub fn total_chunks(&self) -> usize {
    1796
    3
            self.chunks.len()
    1797
    3
        }
    1798
    1799
        /// Get total tokens
    1800
    4
        pub fn total_tokens(&self) -> usize {
    1801
    4
            self.chunks.iter().map(Vec::len).sum()
    1802
    4
        }
    1803
    1804
        /// Check if there are more chunks to process
    1805
    2
        pub fn has_more_chunks(&self) -> bool {
    1806
    2
            self.current_chunk < self.chunks.len()
    1807
    2
        }
    1808
    1809
        /// Get the next chunk to process
    1810
        ///
    1811
        /// Returns None if all chunks are processed
    1812
    35
        pub fn next_chunk(&mut self) -> Option<&[u32]> {
    1813
    35
            if self.start_time.is_none() {
    1814
    4
                self.start_time = Some(std::time::Instant::now());
    1815
    31
            }
    1816
    1817
    35
            if self.current_chunk < self.chunks.len() {
    1818
    31
                let chunk = &self.chunks[self.current_chunk];
    1819
    31
                Some(chunk.as_slice())
    1820
            } else {
    1821
    4
                None
    1822
            }
    1823
    35
        }
    1824
    1825
        /// Mark current chunk as complete
    1826
    31
        pub fn complete_chunk(&mut self, chunk_time_ms: f64) {
    1827
    31
            if self.current_chunk < self.chunks.len() {
    1828
    31
                self.tokens_processed += self.chunks[self.current_chunk].len();
    1829
    31
                self.chunk_times_ms.push(chunk_time_ms);
    1830
    31
                self.current_chunk += 1;
    1831
    31
            }0
    1832
    31
        }
    1833
    1834
        /// Get progress after completing a chunk
    1835
    2
        pub fn progress(&self) -> ChunkProgress {
    1836
    2
            let cumulative_time_ms: f64 = self.chunk_times_ms.iter().sum();
    1837
    1838
    2
            ChunkProgress {
    1839
    2
                chunk_idx: self.current_chunk.saturating_sub(1),
    1840
    2
                total_chunks: self.chunks.len(),
    1841
    2
                tokens_processed: self.tokens_processed,
    1842
    2
                total_tokens: self.total_tokens(),
    1843
    2
                chunk_time_ms: self.chunk_times_ms.last().copied().unwrap_or(0.0),
    1844
    2
                cumulative_time_ms,
    1845
    2
            }
    1846
    2
        }
    1847
    1848
        /// Get estimated time to first token (after first chunk)
    1849
    1
        pub fn estimated_ttft_ms(&self) -> f64 {
    1850
    1
            if let Some(first_chunk_time) = self.chunk_times_ms.first() {
    1851
    1
                *first_chunk_time
    1852
            } else {
    1853
                // Estimate based on chunk size and typical throughput
    1854
    0
                let tokens = self.chunks.first().map_or(0, Vec::len);
    1855
                // Conservative estimate: 0.5ms per token for prefill
    1856
    0
                tokens as f64 * 0.5
    1857
            }
    1858
    1
        }
    1859
    1860
        /// Get statistics after completion
    1861
    1
        pub fn stats(&self) -> ChunkedPrefillStats {
    1862
    1
            let total_time_ms: f64 = self.chunk_times_ms.iter().sum();
    1863
    1
            let total_tokens = self.total_tokens();
    1864
    1
            let avg_chunk_time_ms = if !self.chunk_times_ms.is_empty() {
    1865
    1
                total_time_ms / self.chunk_times_ms.len() as f64
    1866
            } else {
    1867
    0
                0.0
    1868
            };
    1869
    1870
            ChunkedPrefillStats {
    1871
    1
                total_chunks: self.chunks.len(),
    1872
    1
                chunk_size: self.config.chunk_size,
    1873
    1
                total_tokens,
    1874
    1
                total_time_ms,
    1875
    1
                avg_chunk_time_ms,
    1876
    1
                ttft_ms: self.estimated_ttft_ms(),
    1877
    1
                tokens_per_second: if total_time_ms > 0.0 {
    1878
    1
                    total_tokens as f64 / (total_time_ms / 1000.0)
    1879
                } else {
    1880
    0
                    0.0
    1881
                },
    1882
            }
    1883
    1
        }
    1884
    }
    1885
    1886
    /// Statistics for chunked prefill
    1887
    #[cfg(feature = "gpu")]
    1888
    #[derive(Debug, Clone)]
    1889
    pub struct ChunkedPrefillStats {
    1890
        /// Total chunks processed
    1891
        pub total_chunks: usize,
    1892
        /// Chunk size used
    1893
        pub chunk_size: usize,
    1894
        /// Total tokens processed
    1895
        pub total_tokens: usize,
    1896
        /// Total time (ms)
    1897
        pub total_time_ms: f64,
    1898
        /// Average time per chunk (ms)
    1899
        pub avg_chunk_time_ms: f64,
    1900
        /// Time to first token (ms)
    1901
        pub ttft_ms: f64,
    1902
        /// Prefill throughput (tokens/sec)
    1903
        pub tokens_per_second: f64,
    1904
    }