## 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/inference/generation.rs

    Line
    Count
    Source
    1
    //! Token generation for OwnedQuantizedModel
    2
    //!
    3
    //! Contains generate, generate_with_cache, generate_with_cache_streaming,
    4
    //! generate_with_scratch, and sampling methods.
    5
    6
    use crate::error::{RealizarError, Result};
    7
    use crate::gguf::ops;
    8
    use crate::gguf::{
    9
        InferenceScratchBuffer, OwnedQuantizedKVCache, OwnedQuantizedModel, QuantizedGenerateConfig,
    10
    };
    11
    #[cfg(feature = "gpu")]
    12
    use crate::gguf::DispatchMetrics;
    13
    use rand::Rng;
    14
    15
    impl OwnedQuantizedModel {
    16
        /// Get most likely next token
    17
        ///
    18
        /// # Errors
    19
        ///
    20
        /// Returns error if forward pass fails
    21
    0
        pub fn predict_next(&self, token_ids: &[u32]) -> Result<u32> {
    22
    0
            let logits = self.forward(token_ids)?;
    23
    0
            let (max_idx, _) = logits
    24
    0
                .iter()
    25
    0
                .enumerate()
    26
    0
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    27
    0
                .ok_or_else(|| RealizarError::InvalidShape {
    28
    0
                    reason: "Empty logits".to_string(),
    29
    0
                })?;
    30
    0
            Ok(max_idx as u32)
    31
    0
        }
    32
    33
        /// Generate tokens using fused Q4_K operations (IMP-100)
    34
        ///
    35
        /// This is the HTTP serving entry point for quantized inference.
    36
        ///
    37
        /// # Arguments
    38
        ///
    39
        /// * `prompt` - Initial token IDs
    40
        /// * `config` - Generation configuration
    41
        ///
    42
        /// # Returns
    43
        ///
    44
        /// Generated token sequence including prompt
    45
        ///
    46
        /// # Errors
    47
        ///
    48
        /// Returns error if forward pass fails
    49
    0
        pub fn generate(&self, prompt: &[u32], config: &QuantizedGenerateConfig) -> Result<Vec<u32>> {
    50
    0
            if prompt.is_empty() {
    51
    0
                return Err(RealizarError::InvalidShape {
    52
    0
                    reason: "Prompt cannot be empty".to_string(),
    53
    0
                });
    54
    0
            }
    55
    56
    0
            let mut tokens = prompt.to_vec();
    57
    0
            let max_len = prompt.len() + config.max_tokens;
    58
    59
    0
            for _ in 0..config.max_tokens {
    60
                // Forward pass with fused Q4_K ops (1.37x faster)
    61
    0
                let logits = self.forward(&tokens)?;
    62
    63
                // Sample next token
    64
    0
                let next_token = if config.temperature == 0.0 || config.top_k == 1 {
    65
                    // Greedy decoding
    66
    0
                    Self::argmax(&logits)
    67
                } else {
    68
                    // Temperature + top-k sampling
    69
    0
                    Self::sample_topk(&logits, config.temperature, config.top_k)
    70
                };
    71
    72
                // Check stop condition
    73
    0
                if config.stop_tokens.contains(&next_token) {
    74
    0
                    break;
    75
    0
                }
    76
    77
    0
                tokens.push(next_token);
    78
    79
                // Check max length
    80
    0
                if tokens.len() >= max_len {
    81
    0
                    break;
    82
    0
                }
    83
            }
    84
    85
    0
            Ok(tokens)
    86
    0
        }
    87
    88
        /// Greedy argmax over logits
    89
    0
        pub(crate) fn argmax(logits: &[f32]) -> u32 {
    90
    0
            logits
    91
    0
                .iter()
    92
    0
                .enumerate()
    93
    0
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    94
    0
                .map_or(0, |(idx, _)| idx as u32)
    95
    0
        }
    96
    97
        /// Top-k sampling with temperature
    98
    0
        pub fn sample_topk(logits: &[f32], temperature: f32, top_k: usize) -> u32 {
    99
            // Apply temperature
    100
    0
            let scaled: Vec<f32> = logits.iter().map(|&x| x / temperature).collect();
    101
    102
            // Get top-k indices
    103
    0
            let mut indexed: Vec<(usize, f32)> = scaled.iter().copied().enumerate().collect();
    104
    0
            indexed.sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
    105
    0
            indexed.truncate(top_k);
    106
    107
            // Softmax over top-k
    108
    0
            let max_val = indexed.first().map_or(0.0, |(_, v)| *v);
    109
    0
            let exp_sum: f32 = indexed.iter().map(|(_, v)| (v - max_val).exp()).sum();
    110
    0
            let probs: Vec<(usize, f32)> = indexed
    111
    0
                .iter()
    112
    0
                .map(|(i, v)| (*i, (v - max_val).exp() / exp_sum))
    113
    0
                .collect();
    114
    115
            // Sample from probability distribution with proper randomness
    116
    0
            let mut rng = rand::thread_rng();
    117
    0
            let r: f32 = rng.gen();
    118
    119
    0
            let mut cumulative = 0.0;
    120
    0
            for &(idx, prob) in &probs {
    121
    0
                cumulative += prob;
    122
    0
                if cumulative >= r {
    123
    0
                    return idx as u32;
    124
    0
                }
    125
            }
    126
    127
    0
            probs.last().map_or(0, |(idx, _)| *idx as u32)
    128
    0
        }
    129
    130
        /// Generate tokens using KV cache for efficient autoregressive decoding (IMP-101)
    131
        ///
    132
        /// This is O(n) per token instead of O(n²) due to KV cache reuse.
    133
        ///
    134
        /// # Arguments
    135
        /// * `prompt` - Input token IDs
    136
        /// * `config` - Generation configuration
    137
        ///
    138
        /// # Returns
    139
        /// Generated token sequence including prompt
    140
        ///
    141
        /// # Errors
    142
        /// Returns error if forward pass fails
    143
    16
        pub fn generate_with_cache(
    144
    16
            &self,
    145
    16
            prompt: &[u32],
    146
    16
            config: &QuantizedGenerateConfig,
    147
    16
        ) -> Result<Vec<u32>> {
    148
    16
            if prompt.is_empty() {
    149
    0
                return Err(RealizarError::InvalidShape {
    150
    0
                    reason: "Prompt cannot be empty".to_string(),
    151
    0
                });
    152
    16
            }
    153
    154
    16
            let max_seq_len = prompt.len() + config.max_tokens;
    155
    16
            let mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
    156
    16
            let mut tokens = prompt.to_vec();
    157
    158
            // Process prompt tokens (prefill), keeping the logits from the last position
    159
            // The logits from processing token[n-1] at position n-1 predict token[n]
    160
    16
            let mut logits = Vec::new();
    161
    60
            for (pos, &token_id) in prompt16.iter16().enumerate16() {
    162
    60
                logits = self.forward_single_with_cache(token_id, &mut cache, pos)?0;
    163
            }
    164
    165
            // Generate new tokens
    166
            // First iteration uses logits from prefill, subsequent use logits from forward pass
    167
    120
            for gen_idx in 0..config.max_tokens16 {
    168
                // DEBUG: Print logits info for first generated token
    169
    120
                if gen_idx == 0 && std::env::var("REALIZAR_DEBUG_LOGITS")16.is_ok16() {
    170
    0
                    let sum: f32 = logits.iter().sum();
    171
    0
                    let max_val = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    172
    0
                    let min_val = logits.iter().copied().fold(f32::INFINITY, f32::min);
    173
    0
                    let top_5: Vec<(usize, f32)> = {
    174
    0
                        let mut indexed: Vec<_> =
    175
    0
                            logits.iter().enumerate().map(|(i, &v)| (i, v)).collect();
    176
    0
                        indexed.sort_by(|(_, a), (_, b)| {
    177
    0
                            b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal)
    178
    0
                        });
    179
    0
                        indexed.into_iter().take(5).collect()
    180
                    };
    181
    0
                    eprintln!(
    182
    0
                        "[DEBUG-LOGITS] len={}, sum={:.4}, min={:.4}, max={:.4}",
    183
    0
                        logits.len(),
    184
                        sum,
    185
                        min_val,
    186
                        max_val
    187
                    );
    188
    0
                    eprintln!("[DEBUG-LOGITS] top 5 token ids and logits: {:?}", top_5);
    189
    0
                    eprintln!(
    190
    0
                        "[DEBUG-LOGITS] logits[0..5]: {:?}",
    191
    0
                        &logits[..5.min(logits.len())]
    192
                    );
    193
    120
                }
    194
    195
                // Sample next token
    196
    120
                let next_token = if config.temperature == 0.0 || config.top_k == 10 {
    197
    120
                    ops::argmax(&logits)
    198
                } else {
    199
    0
                    crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    200
                };
    201
    202
                // DEBUG: Print selected token
    203
    120
                if gen_idx == 0 && std::env::var("REALIZAR_DEBUG_LOGITS")16.is_ok16() {
    204
    0
                    eprintln!(
    205
    0
                        "[DEBUG-LOGITS] selected token: {} (logit={:.4})",
    206
    0
                        next_token,
    207
    0
                        logits.get(next_token as usize).copied().unwrap_or(f32::NAN)
    208
    0
                    );
    209
    120
                }
    210
    211
                // Check stop condition
    212
    120
                if config.stop_tokens.contains(&next_token) {
    213
    0
                    break;
    214
    120
                }
    215
    216
    120
                tokens.push(next_token);
    217
    218
                // Check max length
    219
    120
                if tokens.len() >= max_seq_len {
    220
    16
                    break;
    221
    104
                }
    222
    223
                // Get logits for next iteration by forwarding the newly sampled token
    224
                // Position is prompt.len() + gen_idx (where token was just added)
    225
    104
                let position = prompt.len() + gen_idx;
    226
    104
                logits = self.forward_single_with_cache(next_token, &mut cache, position)?0;
    227
            }
    228
    229
    16
            Ok(tokens)
    230
    16
        }
    231
    232
        /// Generate tokens with streaming callback (PMAT-087)
    233
        ///
    234
        /// Same as `generate_with_cache` but calls `on_token` after each token
    235
        /// is generated, enabling true streaming to clients.
    236
        ///
    237
        /// # Arguments
    238
        /// * `prompt` - Input token IDs
    239
        /// * `config` - Generation configuration
    240
        /// * `on_token` - Callback called for each generated token. Return `false` to stop.
    241
        ///
    242
        /// # Returns
    243
        /// Generated token sequence including prompt
    244
        ///
    245
        /// # Errors
    246
        /// Returns error if generation fails
    247
    0
        pub fn generate_with_cache_streaming<F>(
    248
    0
            &self,
    249
    0
            prompt: &[u32],
    250
    0
            config: &QuantizedGenerateConfig,
    251
    0
            mut on_token: F,
    252
    0
        ) -> Result<Vec<u32>>
    253
    0
        where
    254
    0
            F: FnMut(u32) -> bool,
    255
        {
    256
    0
            if prompt.is_empty() {
    257
    0
                return Err(RealizarError::InvalidShape {
    258
    0
                    reason: "Prompt cannot be empty".to_string(),
    259
    0
                });
    260
    0
            }
    261
    262
    0
            let max_seq_len = prompt.len() + config.max_tokens;
    263
    0
            let mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
    264
    0
            let mut tokens = prompt.to_vec();
    265
    266
            // Process prompt tokens (prefill)
    267
    0
            let mut logits = Vec::new();
    268
    0
            for (pos, &token_id) in prompt.iter().enumerate() {
    269
    0
                logits = self.forward_single_with_cache(token_id, &mut cache, pos)?;
    270
            }
    271
    272
            // Generate new tokens with streaming
    273
    0
            for gen_idx in 0..config.max_tokens {
    274
                // Sample next token
    275
    0
                let next_token = if config.temperature == 0.0 || config.top_k == 1 {
    276
    0
                    ops::argmax(&logits)
    277
                } else {
    278
    0
                    crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    279
                };
    280
    281
                // Check stop condition
    282
    0
                if config.stop_tokens.contains(&next_token) {
    283
    0
                    break;
    284
    0
                }
    285
    286
    0
                tokens.push(next_token);
    287
    288
                // PMAT-087: Call streaming callback - stop if it returns false
    289
    0
                if !on_token(next_token) {
    290
    0
                    break;
    291
    0
                }
    292
    293
                // Check max length
    294
    0
                if tokens.len() >= max_seq_len {
    295
    0
                    break;
    296
    0
                }
    297
    298
                // Get logits for next iteration
    299
    0
                let position = prompt.len() + gen_idx;
    300
    0
                logits = self.forward_single_with_cache(next_token, &mut cache, position)?;
    301
            }
    302
    303
    0
            Ok(tokens)
    304
    0
        }
    305
    306
        /// Generate tokens with zero-allocation inference (IMP-131)
    307
        ///
    308
        /// This is the highest-performance generation path. Uses pre-allocated
    309
        /// scratch buffers to eliminate per-token allocations, providing ~3-4x
    310
        /// speedup over allocating variants.
    311
        ///
    312
        /// Performance characteristics:
    313
        /// - Single allocation at start (scratch buffer + KV cache)
    314
        /// - Zero allocations per generated token
    315
        /// - ~500KB saved per token for TinyLlama-1.1B
    316
        ///
    317
        /// # Arguments
    318
        /// * `prompt` - Input token IDs
    319
        /// * `config` - Generation configuration
    320
        ///
    321
        /// # Returns
    322
        /// Generated token sequence including prompt
    323
        ///
    324
        /// # Errors
    325
        /// Returns error if forward pass fails
    326
    0
        pub fn generate_with_scratch(
    327
    0
            &self,
    328
    0
            prompt: &[u32],
    329
    0
            config: &QuantizedGenerateConfig,
    330
    0
        ) -> Result<Vec<u32>> {
    331
    0
            if prompt.is_empty() {
    332
    0
                return Err(RealizarError::InvalidShape {
    333
    0
                    reason: "Prompt cannot be empty".to_string(),
    334
    0
                });
    335
    0
            }
    336
    337
    0
            let max_seq_len = prompt.len() + config.max_tokens;
    338
    0
            let mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
    339
    0
            let mut scratch = InferenceScratchBuffer::from_config(&self.config);
    340
    0
            let mut tokens = prompt.to_vec();
    341
    342
            // Process prompt tokens (prefill) - uses scratch buffers
    343
    0
            for (pos, &token_id) in prompt.iter().enumerate() {
    344
    0
                self.forward_single_with_scratch(token_id, &mut cache, pos, &mut scratch)?;
    345
            }
    346
    347
            // Generate new tokens - zero allocations per token
    348
            // PAR-126: Fixed loop structure to match generate_with_cache:
    349
            // 1. Sample from current logits (prefill on first iter, previous forward otherwise)
    350
            // 2. Then run forward on the new token to get logits for next iteration
    351
    0
            for gen_idx in 0..config.max_tokens {
    352
                // Sample next token from current logits (prefill logits on first iter)
    353
    0
                let next_token = if config.temperature == 0.0 || config.top_k == 1 {
    354
    0
                    ops::argmax(&scratch.logits)
    355
                } else {
    356
    0
                    crate::gguf::OwnedQuantizedModel::sample_topk(&scratch.logits, config.temperature, config.top_k)
    357
                };
    358
    359
                // Check stop condition
    360
    0
                if config.stop_tokens.contains(&next_token) {
    361
    0
                    break;
    362
    0
                }
    363
    364
    0
                tokens.push(next_token);
    365
    366
                // Check max length
    367
    0
                if tokens.len() >= max_seq_len {
    368
    0
                    break;
    369
    0
                }
    370
    371
                // Get logits for next iteration by forwarding the new token
    372
    0
                let position = prompt.len() + gen_idx;
    373
    0
                self.forward_single_with_scratch(next_token, &mut cache, position, &mut scratch)?;
    374
            }
    375
    376
    0
            Ok(tokens)
    377
    0
        }
    378
    379
        /// Generate tokens with adaptive CPU/GPU attention (IMP-125)
    380
        ///
    381
        /// This variant of `generate_with_cache` uses `forward_single_with_cache_adaptive`
    382
        /// to automatically select between CPU and GPU backends based on cache length.
    383
        /// It also records dispatch decisions to the provided metrics tracker.
    384
        ///
    385
        /// # Arguments
    386
        /// * `prompt` - Initial token IDs
    387
        /// * `config` - Generation configuration
    388
        /// * `metrics` - Dispatch metrics tracker for CPU/GPU decision recording
    389
        ///
    390
        /// # Returns
    391
        /// Generated token sequence including prompt
    392
        ///
    393
        /// # Errors
    394
        /// Returns error if forward pass fails
    395
        #[cfg(feature = "gpu")]
    396
    9
        pub fn generate_with_cache_adaptive(
    397
    9
            &self,
    398
    9
            prompt: &[u32],
    399
    9
            config: &QuantizedGenerateConfig,
    400
    9
            metrics: &std::sync::Arc<DispatchMetrics>,
    401
    9
        ) -> Result<Vec<u32>> {
    402
    9
            if prompt.is_empty() {
    403
    0
                return Err(RealizarError::InvalidShape {
    404
    0
                    reason: "Prompt cannot be empty".to_string(),
    405
    0
                });
    406
    9
            }
    407
    408
    9
            let max_seq_len = prompt.len() + config.max_tokens;
    409
    9
            let mut cache = OwnedQuantizedKVCache::from_config(&self.config, max_seq_len);
    410
    9
            let mut tokens = prompt.to_vec();
    411
    412
            // Process prompt tokens (prefill) with adaptive attention
    413
            // Keep the logits from the last position for the first generated token
    414
    9
            let mut logits = Vec::new();
    415
    33
            for (pos, &token_id) in prompt9.iter9().enumerate9() {
    416
    33
                logits = self.forward_single_with_cache_adaptive(token_id, &mut cache, pos, metrics)?0;
    417
            }
    418
    419
            // Generate new tokens with adaptive attention
    420
    226
            for gen_idx in 0..config.max_tokens9 {
    421
                // Sample next token from current logits
    422
    226
                let next_token = if config.temperature == 0.0 || config.top_k == 10 {
    423
    226
                    ops::argmax(&logits)
    424
                } else {
    425
    0
                    crate::gguf::OwnedQuantizedModel::sample_topk(&logits, config.temperature, config.top_k)
    426
                };
    427
    428
                // Check stop condition
    429
    226
                if config.stop_tokens.contains(&next_token) {
    430
    0
                    break;
    431
    226
                }
    432
    433
    226
                tokens.push(next_token);
    434
    435
                // Check max length
    436
    226
                if tokens.len() >= max_seq_len {
    437
    9
                    break;
    438
    217
                }
    439
    440
                // Get logits for next iteration by forwarding the newly sampled token
    441
    217
                let position = prompt.len() + gen_idx;
    442
    217
                logits =
    443
    217
                    self.forward_single_with_cache_adaptive(next_token, &mut cache, position, metrics)?0;
    444
            }
    445
    446
    9
            Ok(tokens)
    447
    9
        }
    448
    }