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

    Line
    Count
    Source
    1
    //! Quantized matrix operations for OwnedQuantizedModel
    2
    //!
    3
    //! Contains embed, fused_matmul, qkv_matmul methods with real implementations
    4
    //! for Q4_0, Q8_0, Q4_K, Q5_K, Q6_K quantization formats.
    5
    6
    use crate::error::{RealizarError, Result};
    7
    use crate::gguf::types::{
    8
        GGUF_TYPE_Q4_0, GGUF_TYPE_Q4_1, GGUF_TYPE_Q4_K, GGUF_TYPE_Q5_0, GGUF_TYPE_Q5_K,
    9
        GGUF_TYPE_Q6_K, GGUF_TYPE_Q8_0,
    10
    };
    11
    use crate::gguf::{ops, OwnedQKVWeights, OwnedQuantizedModel, OwnedQuantizedTensor};
    12
    13
    impl OwnedQuantizedModel {
    14
        /// Look up token embeddings (public for debugging PAR-001)
    15
    718
        pub fn embed(&self, token_ids: &[u32]) -> Vec<f32> {
    16
    718
            let hidden_dim = self.config.hidden_dim;
    17
    718
            let mut embeddings = Vec::with_capacity(token_ids.len() * hidden_dim);
    18
    19
    1.48k
            for &token_id764 in token_ids {
    20
    764
                let start = (token_id as usize) * hidden_dim;
    21
    764
                let end = start + hidden_dim;
    22
    764
                if end <= self.token_embedding.len() {
    23
    763
                    embeddings.extend_from_slice(&self.token_embedding[start..end]);
    24
    763
                } else {
    25
    1
                    embeddings.extend(std::iter::repeat_n(0.0, hidden_dim));
    26
    1
                }
    27
            }
    28
    29
    718
            embeddings
    30
    718
        }
    31
    32
        /// Look up single token embedding into pre-allocated buffer (IMP-131)
    33
    2
        pub(crate) fn embed_into(&self, token_id: u32, output: &mut [f32]) {
    34
    2
            let hidden_dim = self.config.hidden_dim;
    35
    2
            let start = (token_id as usize) * hidden_dim;
    36
    2
            let end = start + hidden_dim;
    37
    2
            if end <= self.token_embedding.len() {
    38
    1
                output[..hidden_dim].copy_from_slice(&self.token_embedding[start..end]);
    39
    1
            } else {
    40
    1
                output[..hidden_dim].iter_mut().for_each(|x| *x = 0.0);
    41
            }
    42
    2
        }
    43
    44
        /// Fused dequantize + matmul for quantized weights
    45
        ///
    46
        /// Supports Q4_0, Q8_0, Q4_1, Q5_0, Q4_K, Q5_K, Q6_K quantization formats.
    47
        /// Uses SIMD-accelerated implementations for optimal performance.
    48
    3.50k
        pub(crate) fn fused_matmul(
    49
    3.50k
            &self,
    50
    3.50k
            input: &[f32],
    51
    3.50k
            weight: &OwnedQuantizedTensor,
    52
    3.50k
        ) -> Result<Vec<f32>> {
    53
            use crate::quantize::{
    54
                dequantize_q4_1, dequantize_q5_0, fused_q4_0_q8_0_parallel_matvec,
    55
                fused_q4k_parallel_matvec, fused_q5k_parallel_matvec, fused_q6k_parallel_matvec,
    56
                fused_q8_0_q8_0_parallel_matvec,
    57
            };
    58
            use trueno::{Matrix as TruenoMatrix, Vector as TruenoVector};
    59
    60
    3.50k
            let in_dim = weight.in_dim;
    61
    3.50k
            let out_dim = weight.out_dim;
    62
    3.50k
            let seq_len = input.len() / in_dim;
    63
    64
            // CUDA path when enabled
    65
            #[cfg(feature = "cuda")]
    66
            if let Some(ref executor_mutex) = self.cuda_executor {
    67
                return self.fused_matmul_cuda(input, weight, executor_mutex);
    68
            }
    69
    70
            // CPU path: For Q4_0, use fused Q8_0 integer SIMD matmul (llama.cpp parity)
    71
    3.50k
            if weight.qtype == GGUF_TYPE_Q4_0 {
    72
    3
                if seq_len == 1 {
    73
    2
                    return fused_q4_0_q8_0_parallel_matvec(&weight.data, input, in_dim, out_dim);
    74
    1
                }
    75
    1
                let mut output = Vec::with_capacity(seq_len * out_dim);
    76
    4
                for s in 0..seq_len1 {
    77
    4
                    let x = &input[s * in_dim..(s + 1) * in_dim];
    78
    4
                    let row_output = fused_q4_0_q8_0_parallel_matvec(&weight.data, x, in_dim, out_dim)?0;
    79
    4
                    output.extend_from_slice(&row_output);
    80
                }
    81
    1
                return Ok(output);
    82
    3.50k
            }
    83
    84
            // CPU path: For Q8_0, use fused Q8_0 × Q8_0 integer SIMD matmul
    85
    3.50k
            if weight.qtype == GGUF_TYPE_Q8_0 {
    86
    3
                if seq_len == 1 {
    87
    2
                    return fused_q8_0_q8_0_parallel_matvec(&weight.data, input, in_dim, out_dim);
    88
    1
                }
    89
    1
                let mut output = Vec::with_capacity(seq_len * out_dim);
    90
    3
                for s in 0..seq_len1 {
    91
    3
                    let x = &input[s * in_dim..(s + 1) * in_dim];
    92
    3
                    let row_output = fused_q8_0_q8_0_parallel_matvec(&weight.data, x, in_dim, out_dim)?0;
    93
    3
                    output.extend_from_slice(&row_output);
    94
                }
    95
    1
                return Ok(output);
    96
    3.49k
            }
    97
    98
            // CPU path: For Q4_1, use dequantize + SIMD matmul
    99
    3.49k
            if weight.qtype == GGUF_TYPE_Q4_1 {
    100
    3
                let weights_f32 = dequantize_q4_1(&weight.data)?0;
    101
    102
    3
                let weight_matrix = match TruenoMatrix::from_vec(out_dim, in_dim, weights_f32) {
    103
    3
                    Ok(m) => m,
    104
                    Err(_) => {
    105
    0
                        return Err(RealizarError::InvalidShape {
    106
    0
                            reason: "Failed to create weight matrix for Q4_1".to_string(),
    107
    0
                        });
    108
                    }
    109
                };
    110
    111
    3
                let mut output = Vec::with_capacity(seq_len * out_dim);
    112
    4
                for s in 0..seq_len3 {
    113
    4
                    let x = &input[s * in_dim..(s + 1) * in_dim];
    114
    4
                    let x_vec = TruenoVector::from_slice(x);
    115
    4
                    match weight_matrix.matvec(&x_vec) {
    116
    4
                        Ok(r) => output.extend_from_slice(r.as_slice()),
    117
                        Err(_) => {
    118
    0
                            return Err(RealizarError::InvalidShape {
    119
    0
                                reason: "SIMD matvec failed for Q4_1".to_string(),
    120
    0
                            });
    121
                        }
    122
                    }
    123
                }
    124
    3
                return Ok(output);
    125
    3.49k
            }
    126
    127
            // CPU path: For Q5_0, use dequantize + SIMD matmul
    128
    3.49k
            if weight.qtype == GGUF_TYPE_Q5_0 {
    129
    3
                let weights_f32 = dequantize_q5_0(&weight.data)?0;
    130
    131
    3
                let weight_matrix = match TruenoMatrix::from_vec(out_dim, in_dim, weights_f32) {
    132
    3
                    Ok(m) => m,
    133
                    Err(_) => {
    134
    0
                        return Err(RealizarError::InvalidShape {
    135
    0
                            reason: "Failed to create weight matrix for Q5_0".to_string(),
    136
    0
                        });
    137
                    }
    138
                };
    139
    140
    3
                let mut output = Vec::with_capacity(seq_len * out_dim);
    141
    4
                for s in 0..seq_len3 {
    142
    4
                    let x = &input[s * in_dim..(s + 1) * in_dim];
    143
    4
                    let x_vec = TruenoVector::from_slice(x);
    144
    4
                    match weight_matrix.matvec(&x_vec) {
    145
    4
                        Ok(r) => output.extend_from_slice(r.as_slice()),
    146
                        Err(_) => {
    147
    0
                            return Err(RealizarError::InvalidShape {
    148
    0
                                reason: "SIMD matvec failed for Q5_0".to_string(),
    149
    0
                            });
    150
                        }
    151
                    }
    152
                }
    153
    3
                return Ok(output);
    154
    3.49k
            }
    155
    156
            // CPU path: Process each position in sequence for Q4_K, Q5_K, Q6_K
    157
    3.49k
            if seq_len > 1 {
    158
    18
                let mut output = Vec::with_capacity(seq_len * out_dim);
    159
    70
                for s in 0..seq_len18 {
    160
    70
                    let x = &input[s * in_dim..(s + 1) * in_dim];
    161
    70
                    let row_output = match weight.qtype {
    162
                        GGUF_TYPE_Q4_K => {
    163
    66
                            fused_q4k_parallel_matvec(&weight.data, x, in_dim, out_dim)?0
    164
                        }
    165
                        GGUF_TYPE_Q5_K => {
    166
    2
                            fused_q5k_parallel_matvec(&weight.data, x, in_dim, out_dim)?0
    167
                        }
    168
                        GGUF_TYPE_Q6_K => {
    169
    2
                            fused_q6k_parallel_matvec(&weight.data, x, in_dim, out_dim)?0
    170
                        }
    171
                        _ => {
    172
    0
                            return Err(RealizarError::UnsupportedOperation {
    173
    0
                                operation: "owned_fused_matmul".to_string(),
    174
    0
                                reason: format!(
    175
    0
                                    "Fused matmul only supports Q4_0/Q4_1/Q5_0/Q8_0/Q4_K/Q5_K/Q6_K, got type {}",
    176
    0
                                    weight.qtype
    177
    0
                                ),
    178
    0
                            });
    179
                        }
    180
                    };
    181
    70
                    output.extend_from_slice(&row_output);
    182
                }
    183
    18
                Ok(output)
    184
            } else {
    185
                // Single position - most common case in generation
    186
    3.47k
                match weight.qtype {
    187
    3.46k
                    GGUF_TYPE_Q4_K => fused_q4k_parallel_matvec(&weight.data, input, in_dim, out_dim),
    188
    2
                    GGUF_TYPE_Q5_K => fused_q5k_parallel_matvec(&weight.data, input, in_dim, out_dim),
    189
    2
                    GGUF_TYPE_Q6_K => fused_q6k_parallel_matvec(&weight.data, input, in_dim, out_dim),
    190
    1
                    _ => Err(RealizarError::UnsupportedOperation {
    191
    1
                        operation: "owned_fused_matmul".to_string(),
    192
    1
                        reason: format!(
    193
    1
                            "Fused matmul only supports Q4_0/Q8_0/Q4_K/Q5_K/Q6_K, got type {}",
    194
    1
                            weight.qtype
    195
    1
                        ),
    196
    1
                    }),
    197
                }
    198
            }
    199
    3.50k
        }
    200
    201
        /// CUDA path for fused matmul
    202
        #[cfg(feature = "cuda")]
    203
        fn fused_matmul_cuda(
    204
            &self,
    205
            input: &[f32],
    206
            weight: &OwnedQuantizedTensor,
    207
            executor_mutex: &std::sync::Mutex<crate::cuda::CudaExecutor>,
    208
        ) -> Result<Vec<f32>> {
    209
            use tracing::info_span;
    210
    211
            let in_dim = weight.in_dim;
    212
            let out_dim = weight.out_dim;
    213
            let seq_len = input.len() / in_dim;
    214
            let gemm_start = std::time::Instant::now();
    215
            let mut output = vec![0.0f32; seq_len * out_dim];
    216
    217
            // Use native quantized GEMV kernels for single-token generation
    218
            if seq_len == 1 {
    219
                let cache_key = format!("{}_{:016x}",
    220
                    match weight.qtype {
    221
                        GGUF_TYPE_Q4_K => "q4k",
    222
                        GGUF_TYPE_Q5_K => "q5k",
    223
                        GGUF_TYPE_Q6_K => "q6k",
    224
                        _ => "unknown",
    225
                    },
    226
                    weight.data.as_ptr() as usize
    227
                );
    228
    229
                if weight.qtype == GGUF_TYPE_Q4_K
    230
                    || weight.qtype == GGUF_TYPE_Q5_K
    231
                    || weight.qtype == GGUF_TYPE_Q6_K
    232
                {
    233
                    let mut executor =
    234
                        executor_mutex
    235
                            .lock()
    236
                            .map_err(|e| RealizarError::UnsupportedOperation {
    237
                                operation: "cuda_lock".to_string(),
    238
                                reason: format!("Failed to acquire CUDA executor lock: {e}"),
    239
                            })?;
    240
    241
                    executor
    242
                        .make_current()
    243
                        .map_err(|e| RealizarError::UnsupportedOperation {
    244
                            operation: "cuda_make_current".to_string(),
    245
                            reason: format!("Failed to set CUDA context current: {e}"),
    246
                        })?;
    247
    248
                    if !executor.has_quantized_weights(&cache_key) {
    249
                        executor
    250
                            .load_quantized_weights(&cache_key, &weight.data)
    251
                            .map_err(|e| RealizarError::UnsupportedOperation {
    252
                                operation: "cuda_cache".to_string(),
    253
                                reason: format!("Failed to cache weights: {e}"),
    254
                            })?;
    255
                    }
    256
    257
                    let result = match weight.qtype {
    258
                        GGUF_TYPE_Q4_K => executor.q4k_gemv_cached(
    259
                            &cache_key,
    260
                            input,
    261
                            &mut output,
    262
                            out_dim as u32,
    263
                            in_dim as u32,
    264
                        ),
    265
                        GGUF_TYPE_Q5_K => executor.q5k_gemv_cached(
    266
                            &cache_key,
    267
                            input,
    268
                            &mut output,
    269
                            out_dim as u32,
    270
                            in_dim as u32,
    271
                        ),
    272
                        GGUF_TYPE_Q6_K => executor.q6k_gemv_cached(
    273
                            &cache_key,
    274
                            input,
    275
                            &mut output,
    276
                            out_dim as u32,
    277
                            in_dim as u32,
    278
                        ),
    279
                        _ => unreachable!(),
    280
                    };
    281
    282
                    result.map_err(|e| RealizarError::UnsupportedOperation {
    283
                        operation: "cuda_gemv".to_string(),
    284
                        reason: format!("CUDA GEMV failed: {e}"),
    285
                    })?;
    286
    287
                    let gemm_duration_us = gemm_start.elapsed().as_micros() as u64;
    288
                    let _span = info_span!(
    289
                        "gpu_kernel:gemv",
    290
                        gpu.backend = "cuda",
    291
                        gpu.dimensions.n = out_dim,
    292
                        gpu.dimensions.k = in_dim,
    293
                        duration_us = gemm_duration_us,
    294
                    )
    295
                    .entered();
    296
    297
                    self.cuda_kernel_count
    298
                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    299
    300
                    return Ok(output);
    301
                }
    302
            }
    303
    304
            // Fallback: Dequantize and use FP32 GEMM
    305
            let dequant_weight = self.dequantize_weight_for_cuda(weight)?;
    306
    307
            {
    308
                let mut executor =
    309
                    executor_mutex
    310
                        .lock()
    311
                        .map_err(|e| RealizarError::UnsupportedOperation {
    312
                            operation: "cuda_gemm_lock".to_string(),
    313
                            reason: format!("Failed to acquire CUDA executor lock: {e}"),
    314
                        })?;
    315
    316
                executor
    317
                    .make_current()
    318
                    .map_err(|e| RealizarError::UnsupportedOperation {
    319
                        operation: "cuda_make_current".to_string(),
    320
                        reason: format!("Failed to set CUDA context current: {e}"),
    321
                    })?;
    322
    323
                executor
    324
                    .gemm(
    325
                        input,
    326
                        &dequant_weight,
    327
                        &mut output,
    328
                        seq_len as u32,
    329
                        out_dim as u32,
    330
                        in_dim as u32,
    331
                    )
    332
                    .map_err(|e| RealizarError::UnsupportedOperation {
    333
                        operation: "cuda_gemm".to_string(),
    334
                        reason: format!("CUDA GEMM failed: {e}"),
    335
                    })?;
    336
            }
    337
    338
            self.cuda_kernel_count
    339
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    340
    341
            Ok(output)
    342
        }
    343
    344
        /// Fused matmul into pre-allocated output buffer
    345
    0
        pub(crate) fn fused_matmul_into(
    346
    0
            &self,
    347
    0
            input: &[f32],
    348
    0
            weight: &OwnedQuantizedTensor,
    349
    0
            output: &mut [f32],
    350
    0
        ) -> Result<()> {
    351
            use crate::quantize::{
    352
                fused_q4_0_q8_0_parallel_matvec_into, fused_q4k_parallel_matvec_into,
    353
                fused_q5k_parallel_matvec_into, fused_q6k_parallel_matvec_into,
    354
                fused_q8_0_q8_0_parallel_matvec_into,
    355
            };
    356
    357
    0
            let in_dim = weight.in_dim;
    358
    0
            let out_dim = weight.out_dim;
    359
    0
            let seq_len = input.len() / in_dim;
    360
    361
            // Only support single-token case for now (most common in generation)
    362
    0
            if seq_len != 1 {
    363
    0
                let result = self.fused_matmul(input, weight)?;
    364
    0
                output[..result.len()].copy_from_slice(&result);
    365
    0
                return Ok(());
    366
    0
            }
    367
    368
    0
            debug_assert!(
    369
    0
                output.len() >= out_dim,
    370
    0
                "Output buffer too small: {} < {}",
    371
    0
                output.len(),
    372
                out_dim
    373
            );
    374
    375
    0
            match weight.qtype {
    376
                GGUF_TYPE_Q4_0 => {
    377
    0
                    fused_q4_0_q8_0_parallel_matvec_into(
    378
    0
                        &weight.data,
    379
    0
                        input,
    380
    0
                        in_dim,
    381
    0
                        &mut output[..out_dim],
    382
                    )
    383
                }
    384
    0
                GGUF_TYPE_Q8_0 => fused_q8_0_q8_0_parallel_matvec_into(
    385
    0
                    &weight.data,
    386
    0
                    input,
    387
    0
                    in_dim,
    388
    0
                    out_dim,
    389
    0
                    &mut output[..out_dim],
    390
                ),
    391
    0
                GGUF_TYPE_Q4_K => fused_q4k_parallel_matvec_into(
    392
    0
                    &weight.data,
    393
    0
                    input,
    394
    0
                    in_dim,
    395
    0
                    out_dim,
    396
    0
                    &mut output[..out_dim],
    397
                ),
    398
    0
                GGUF_TYPE_Q5_K => fused_q5k_parallel_matvec_into(
    399
    0
                    &weight.data,
    400
    0
                    input,
    401
    0
                    in_dim,
    402
    0
                    out_dim,
    403
    0
                    &mut output[..out_dim],
    404
                ),
    405
    0
                GGUF_TYPE_Q6_K => fused_q6k_parallel_matvec_into(
    406
    0
                    &weight.data,
    407
    0
                    input,
    408
    0
                    in_dim,
    409
    0
                    out_dim,
    410
    0
                    &mut output[..out_dim],
    411
                ),
    412
                _ => {
    413
    0
                    let result = self.fused_matmul(input, weight)?;
    414
    0
                    output[..result.len()].copy_from_slice(&result);
    415
    0
                    Ok(())
    416
                }
    417
            }
    418
    0
        }
    419
    420
        /// QKV projection matmul
    421
    714
        pub fn qkv_matmul(&self, input: &[f32], qkv: &OwnedQKVWeights) -> Result<Vec<f32>> {
    422
    714
            let hidden_dim = self.config.hidden_dim;
    423
    714
            match qkv {
    424
    714
                OwnedQKVWeights::Fused(ref weight) => self.fused_matmul(input, weight),
    425
                OwnedQKVWeights::Separate {
    426
    0
                    ref q,
    427
    0
                    ref k,
    428
    0
                    ref v,
    429
                } => {
    430
    0
                    let seq_len = input.len() / hidden_dim;
    431
    432
    0
                    let q_out = self.fused_matmul(input, q)?;
    433
    0
                    let k_out = self.fused_matmul(input, k)?;
    434
    0
                    let v_out = self.fused_matmul(input, v)?;
    435
    436
                    // Interleave Q, K, V for each position
    437
    0
                    let qkv_dim = q.out_dim + k.out_dim + v.out_dim;
    438
    0
                    let mut output = Vec::with_capacity(seq_len * qkv_dim);
    439
    0
                    for s in 0..seq_len {
    440
    0
                        output.extend_from_slice(&q_out[s * q.out_dim..(s + 1) * q.out_dim]);
    441
    0
                        output.extend_from_slice(&k_out[s * k.out_dim..(s + 1) * k.out_dim]);
    442
    0
                        output.extend_from_slice(&v_out[s * v.out_dim..(s + 1) * v.out_dim]);
    443
    0
                    }
    444
    0
                    Ok(output)
    445
                }
    446
            }
    447
    714
        }
    448
    449
        /// QKV matmul into pre-allocated buffer
    450
    0
        pub fn qkv_matmul_into(
    451
    0
            &self,
    452
    0
            input: &[f32],
    453
    0
            qkv: &OwnedQKVWeights,
    454
    0
            output: &mut [f32],
    455
    0
        ) -> Result<()> {
    456
    0
            match qkv {
    457
    0
                OwnedQKVWeights::Fused(ref weight) => self.fused_matmul_into(input, weight, output),
    458
                OwnedQKVWeights::Separate {
    459
    0
                    ref q,
    460
    0
                    ref k,
    461
    0
                    ref v,
    462
                } => {
    463
    0
                    let q_dim = q.out_dim;
    464
    0
                    let k_dim = k.out_dim;
    465
    0
                    let v_dim = v.out_dim;
    466
    467
    0
                    self.fused_matmul_into(input, q, &mut output[..q_dim])?;
    468
    0
                    self.fused_matmul_into(input, k, &mut output[q_dim..q_dim + k_dim])?;
    469
    0
                    self.fused_matmul_into(
    470
    0
                        input,
    471
    0
                        v,
    472
    0
                        &mut output[q_dim + k_dim..q_dim + k_dim + v_dim],
    473
    0
                    )?;
    474
    475
    0
                    Ok(())
    476
                }
    477
            }
    478
    0
        }
    479
    480
        /// Layer normalization
    481
    18
        pub fn layer_norm(
    482
    18
            &self,
    483
    18
            input: &[f32],
    484
    18
            weight: &[f32],
    485
    18
            bias: Option<&[f32]>,
    486
    18
            eps: f32,
    487
    18
        ) -> Vec<f32> {
    488
    18
            ops::layer_norm(input, weight, bias, eps)
    489
    18
        }
    490
    491
        /// Add bias to activations
    492
    0
        pub fn add_bias(&self, input: &mut [f32], bias: &[f32]) {
    493
    0
            for (x, b) in input.iter_mut().zip(bias.iter()) {
    494
    0
                *x += b;
    495
    0
            }
    496
    0
        }
    497
    498
        /// GELU activation
    499
    6
        pub fn gelu(&self, input: &mut [f32]) {
    500
    2.30k
            for x in input6.iter_mut6() {
    501
    2.30k
                *x = 0.5 * *x * (1.0 + (*x * 0.7978845608 * (1.0 + 0.044715 * *x * *x)).tanh());
    502
    2.30k
            }
    503
    6
        }
    504
    505
        /// Fused RMSNorm + matmul helper
    506
    0
        fn fused_rmsnorm_matmul(
    507
    0
            &self,
    508
    0
            input: &[f32],
    509
    0
            norm_weight: &[f32],
    510
    0
            eps: f32,
    511
    0
            weight: &OwnedQuantizedTensor,
    512
    0
        ) -> Result<Vec<f32>> {
    513
            use crate::quantize::fused_rmsnorm_q4_0_matmul;
    514
    515
            // Only use fused path for Q4_0 weights (most common)
    516
    0
            if weight.qtype == GGUF_TYPE_Q4_0 && input.len() == weight.in_dim {
    517
    0
                return fused_rmsnorm_q4_0_matmul(
    518
    0
                    input,
    519
    0
                    norm_weight,
    520
    0
                    eps,
    521
    0
                    &weight.data,
    522
    0
                    weight.in_dim,
    523
    0
                    weight.out_dim,
    524
                );
    525
    0
            }
    526
    527
            // Fallback to separate RMSNorm + matmul for other types
    528
    0
            let normed = ops::rms_norm(input, norm_weight, eps);
    529
    0
            self.fused_matmul(&normed, weight)
    530
    0
        }
    531
    532
        /// Fused RMSNorm + QKV matmul
    533
    0
        pub fn fused_rmsnorm_qkv_matmul(
    534
    0
            &self,
    535
    0
            input: &[f32],
    536
    0
            norm_weight: &[f32],
    537
    0
            eps: f32,
    538
    0
            qkv: &OwnedQKVWeights,
    539
    0
        ) -> Result<Vec<f32>> {
    540
    0
            match qkv {
    541
    0
                OwnedQKVWeights::Fused(ref weight) => {
    542
    0
                    self.fused_rmsnorm_matmul(input, norm_weight, eps, weight)
    543
                }
    544
                OwnedQKVWeights::Separate {
    545
    0
                    ref q,
    546
    0
                    ref k,
    547
    0
                    ref v,
    548
                } => {
    549
                    // For separate Q/K/V, normalize once and reuse
    550
    0
                    let normed = ops::rms_norm(input, norm_weight, eps);
    551
    552
    0
                    let q_out = self.fused_matmul(&normed, q)?;
    553
    0
                    let k_out = self.fused_matmul(&normed, k)?;
    554
    0
                    let v_out = self.fused_matmul(&normed, v)?;
    555
    556
    0
                    let qkv_dim = q.out_dim + k.out_dim + v.out_dim;
    557
    0
                    let mut output = Vec::with_capacity(qkv_dim);
    558
    0
                    output.extend_from_slice(&q_out);
    559
    0
                    output.extend_from_slice(&k_out);
    560
    0
                    output.extend_from_slice(&v_out);
    561
    0
                    Ok(output)
    562
                }
    563
            }
    564
    0
        }
    565
    566
        /// Fused RMSNorm + LM head
    567
    0
        pub fn fused_rmsnorm_lm_head(&self, input: &[f32]) -> Result<Vec<f32>> {
    568
            use crate::quantize::fused_rmsnorm_q4_0_matmul;
    569
    570
            // Only use fused path for Q4_0 weights
    571
    0
            if self.lm_head_weight.qtype == GGUF_TYPE_Q4_0
    572
    0
                && input.len() == self.lm_head_weight.in_dim
    573
            {
    574
    0
                return fused_rmsnorm_q4_0_matmul(
    575
    0
                    input,
    576
    0
                    &self.output_norm_weight,
    577
    0
                    self.config.eps,
    578
    0
                    &self.lm_head_weight.data,
    579
    0
                    self.lm_head_weight.in_dim,
    580
    0
                    self.lm_head_weight.out_dim,
    581
                );
    582
    0
            }
    583
    584
            // Fallback to separate RMSNorm + matmul for other types
    585
    0
            let normed = ops::rms_norm(input, &self.output_norm_weight, self.config.eps);
    586
    0
            self.fused_matmul(&normed, &self.lm_head_weight)
    587
    0
        }
    588
    589
        /// Fused RMSNorm + FFN up/gate projections for SwiGLU
    590
    0
        pub fn fused_rmsnorm_ffn_up_gate(
    591
    0
            &self,
    592
    0
            input: &[f32],
    593
    0
            norm_weight: &[f32],
    594
    0
            eps: f32,
    595
    0
            up_weight: &OwnedQuantizedTensor,
    596
    0
            gate_weight: &OwnedQuantizedTensor,
    597
    0
        ) -> Result<(Vec<f32>, Vec<f32>)> {
    598
            use crate::quantize::fused_rmsnorm_ffn_up_gate;
    599
    600
            // Only use fused path for Q4_0 weights
    601
    0
            if up_weight.qtype == GGUF_TYPE_Q4_0
    602
    0
                && gate_weight.qtype == GGUF_TYPE_Q4_0
    603
    0
                && input.len() == up_weight.in_dim
    604
    0
                && up_weight.in_dim == gate_weight.in_dim
    605
    0
                && up_weight.out_dim == gate_weight.out_dim
    606
            {
    607
    0
                return fused_rmsnorm_ffn_up_gate(
    608
    0
                    input,
    609
    0
                    norm_weight,
    610
    0
                    eps,
    611
    0
                    &up_weight.data,
    612
    0
                    &gate_weight.data,
    613
    0
                    up_weight.in_dim,
    614
    0
                    up_weight.out_dim,
    615
                );
    616
    0
            }
    617
    618
            // Fallback to separate RMSNorm + matmuls for other types
    619
    0
            let normed = ops::rms_norm(input, norm_weight, eps);
    620
    0
            let up_out = self.fused_matmul(&normed, up_weight)?;
    621
    0
            let gate_out = self.fused_matmul(&normed, gate_weight)?;
    622
    0
            Ok((up_out, gate_out))
    623
    0
        }
    624
    625
        /// Q8K QKV matmul into buffer
    626
        ///
    627
        /// Uses pre-quantized Q8K activations for faster matmul with Q4K weights.
    628
        #[allow(unused_variables)]
    629
    0
        pub fn qkv_matmul_q8k_into(
    630
    0
            &self,
    631
    0
            input: &[f32],
    632
    0
            qkv: &OwnedQKVWeights,
    633
    0
            output: &mut [f32],
    634
    0
            scales: &[f32],
    635
    0
            quants: &[i8],
    636
    0
        ) -> Result<()> {
    637
            // Fall back to regular qkv_matmul_into for now
    638
            // TODO: Implement Q8K-accelerated path using pre-quantized activations
    639
    0
            self.qkv_matmul_into(input, qkv, output)
    640
    0
        }
    641
    642
        /// Helper to dequantize weights for CUDA GEMM
    643
        #[cfg(feature = "cuda")]
    644
        fn dequantize_weight_for_cuda(&self, weight: &OwnedQuantizedTensor) -> Result<Vec<f32>> {
    645
            use crate::quantize::{
    646
                dequantize_q4_0, dequantize_q4_1, dequantize_q4_k, dequantize_q5_0, dequantize_q5_k,
    647
                dequantize_q6_k, dequantize_q8_0,
    648
            };
    649
    650
            match weight.qtype {
    651
                GGUF_TYPE_Q4_0 => dequantize_q4_0(&weight.data),
    652
                GGUF_TYPE_Q4_1 => dequantize_q4_1(&weight.data),
    653
                GGUF_TYPE_Q5_0 => dequantize_q5_0(&weight.data),
    654
                GGUF_TYPE_Q8_0 => dequantize_q8_0(&weight.data),
    655
                GGUF_TYPE_Q4_K => dequantize_q4_k(&weight.data),
    656
                GGUF_TYPE_Q5_K => dequantize_q5_k(&weight.data),
    657
                GGUF_TYPE_Q6_K => dequantize_q6_k(&weight.data),
    658
                _ => Err(RealizarError::UnsupportedOperation {
    659
                    operation: "dequantize_weight_for_cuda".to_string(),
    660
                    reason: format!("Unsupported quantization type: {}", weight.qtype),
    661
                }),
    662
            }
    663
        }
    664
    }