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

    Line
    Count
    Source
    1
    //! APR Transformer Format for WASM-compatible LLM inference
    2
    //!
    3
    //! This module provides a WASM-compatible transformer implementation that stores
    4
    //! all weights as F32, enabling fair comparison between APR and GGUF formats.
    5
    //!
    6
    //! ## Design Goals
    7
    //!
    8
    //! 1. **WASM Compatibility**: Pure F32 weights, no SIMD requirements
    9
    //! 2. **Fair Comparison**: Same inference algorithm as GGUFTransformer
    10
    //! 3. **Serialization**: APR format with model type `TransformerLM` (0x0050)
    11
    //!
    12
    //! ## Example
    13
    //!
    14
    //! ```rust,ignore
    15
    //! use realizar::apr_transformer::AprTransformer;
    16
    //! use realizar::gguf::{GGUFModel, GGUFTransformer};
    17
    //!
    18
    //! // Load GGUF model
    19
    //! let gguf_data = std::fs::read("model.gguf")?;
    20
    //! let gguf_model = GGUFModel::from_bytes(&gguf_data)?;
    21
    //! let gguf_transformer = GGUFTransformer::from_gguf(&gguf_model, &gguf_data)?;
    22
    //!
    23
    //! // Convert to APR format
    24
    //! let apr_transformer = AprTransformer::from_gguf_transformer(&gguf_transformer);
    25
    //!
    26
    //! // Run inference (should match GGUF output)
    27
    //! let logits = apr_transformer.forward(&[1, 2, 3, 4])?;
    28
    //! ```
    29
    30
    use std::fs::File;
    31
    use std::path::Path;
    32
    33
    use serde::{Deserialize, Serialize};
    34
    35
    use crate::error::{RealizarError, Result};
    36
    37
    // PMAT-802: Extracted modules
    38
    mod config;
    39
    mod dequant;
    40
    mod helpers;
    41
    mod loader;
    42
    mod q4_simd;
    43
    pub use config::{AprKVCache, GenerateConfig, AprTransformerConfig, AprTransformerLayer, Q4KLayerWeights};
    44
    pub use loader::{MmapAprTransformer, AprQuantizationType, QuantizedAprTransformer, APR_TRANSFORMER_HEADER_SIZE};
    45
    pub use q4_simd::{QuantizedAprTensorQ4, QuantizedAprLayerQ4, QuantizedAprTransformerQ4, AprInferenceScratch};
    46
    use dequant::{dequantize_q4_k_apr, dequantize_q6_k_apr};
    47
    use helpers::{matmul_q4k_rowmajor, matmul_q6k_rowmajor, simd_dot_f32, simd_add_weighted};
    48
    49
    // APR Benchmark Infrastructure (Y6) - extracted from mod.rs (PMAT-802)
    50
    mod benchmark;
    51
    pub use benchmark::{
    52
        AprBenchmarkResult, AprBenchmarkRunner, AprLoadResult, AprParityComparison, AprPrefillResult,
    53
        APR_CPU_DECODE_THRESHOLD_TOK_S, APR_PARITY_THRESHOLD_PCT, APR_PREFILL_THRESHOLD_TOK_S,
    54
    };
    55
    56
    /// APR Transformer model with all weights
    57
    ///
    58
    /// For Q4K models, raw Q4K bytes can be stored in `q4k_layers` to enable
    59
    /// fused kernel inference (F-GPU-130) without full dequantization overhead.
    60
    #[derive(Debug, Clone, Serialize, Deserialize)]
    61
    pub struct AprTransformer {
    62
        /// Model configuration
    63
        pub config: AprTransformerConfig,
    64
        /// Token embedding weights [vocab_size * hidden_dim]
    65
        pub token_embedding: Vec<f32>,
    66
        /// Transformer layers
    67
        pub layers: Vec<AprTransformerLayer>,
    68
        /// Output norm weight [hidden_dim]
    69
        pub output_norm_weight: Vec<f32>,
    70
        /// Output norm bias (optional) [hidden_dim]
    71
        pub output_norm_bias: Option<Vec<f32>>,
    72
        /// LM head weight [hidden_dim * vocab_size]
    73
        pub lm_head_weight: Vec<f32>,
    74
        /// LM head bias (optional) [vocab_size]
    75
        pub lm_head_bias: Option<Vec<f32>>,
    76
        /// Q4K raw layer weights for fused kernel inference (F-GPU-130)
    77
        /// When present, enables direct Q4K matmul without dequantization
    78
        #[serde(default)]
    79
        pub q4k_layers: Option<Vec<Q4KLayerWeights>>,
    80
        /// LM head weight in Q6K format for fused kernel inference
    81
        /// When present, enables direct Q6K matmul without dequantization
    82
        #[serde(default)]
    83
        pub lm_head_weight_q6k: Option<Vec<u8>>,
    84
        /// LM head weight in Q4K format for fused kernel inference
    85
        #[serde(default)]
    86
        pub lm_head_weight_q4k: Option<Vec<u8>>,
    87
    }
    88
    89
    impl AprTransformer {
    90
        /// Load APR transformer from an APR v2 file
    91
        ///
    92
        /// Parses the APR v2 format (magic "APR2") and extracts transformer weights.
    93
        ///
    94
        /// # Arguments
    95
        ///
    96
        /// * `path` - Path to .apr file
    97
        ///
    98
        /// # Returns
    99
        ///
    100
        /// Loaded transformer ready for inference
    101
        ///
    102
        /// # Errors
    103
        ///
    104
        /// Returns error if file cannot be read or parsed
    105
        ///
    106
        /// # Example
    107
        ///
    108
        /// ```rust,ignore
    109
        /// let transformer = AprTransformer::from_apr_file("model.apr")?;
    110
        /// let logits = transformer.forward(&[1, 2, 3])?;
    111
        /// ```
    112
    0
        pub fn from_apr_file<P: AsRef<Path>>(path: P) -> Result<Self> {
    113
            use std::io::Read;
    114
    115
    0
            let mut file = File::open(path.as_ref()).map_err(|e| RealizarError::IoError {
    116
    0
                message: format!("Failed to open APR file: {e}"),
    117
    0
            })?;
    118
    119
    0
            let mut data = Vec::new();
    120
    0
            file.read_to_end(&mut data)
    121
    0
                .map_err(|e| RealizarError::IoError {
    122
    0
                    message: format!("Failed to read APR file: {e}"),
    123
    0
                })?;
    124
    125
    0
            Self::from_apr_bytes(&data)
    126
    0
        }
    127
    128
        /// Load APR transformer from bytes
    129
        ///
    130
        /// Parses APR v2 format from memory buffer.
    131
    0
        pub fn from_apr_bytes(data: &[u8]) -> Result<Self> {
    132
            // Check minimum size for header
    133
    0
            if data.len() < 64 {
    134
    0
                return Err(RealizarError::FormatError {
    135
    0
                    reason: format!("APR file too small: {} bytes (need 64)", data.len()),
    136
    0
                });
    137
    0
            }
    138
    139
            // Check magic - first 3 bytes must be "APR", 4th byte is version (0, '1', or '2')
    140
    0
            let magic = &data[0..4];
    141
    0
            if magic[0..3] != *b"APR"
    142
    0
                || (magic[3] != 0 && magic[3] != b'1' && magic[3] != b'2')
    143
            {
    144
    0
                return Err(RealizarError::FormatError {
    145
    0
                    reason: format!(
    146
    0
                        "Invalid APR magic: {:?}, expected APR followed by version byte",
    147
    0
                        String::from_utf8_lossy(magic)
    148
    0
                    ),
    149
    0
                });
    150
    0
            }
    151
    152
            // Parse header
    153
            // APR header layout:
    154
            //   0-3: Magic "APR\0"
    155
            //   4-5: Version major.minor
    156
            //   6-7: Flags
    157
            //   8-11: Tensor count
    158
            //   12-19: Metadata offset
    159
            //   20-23: Metadata size
    160
            //   24-31: Tensor index offset
    161
            //   32-39: Data offset
    162
            //   40-43: Checksum
    163
            //   44-63: Reserved
    164
    165
    0
            let tensor_count = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
    166
    0
            let metadata_offset = u64::from_le_bytes([
    167
    0
                data[12], data[13], data[14], data[15], data[16], data[17], data[18], data[19],
    168
    0
            ]) as usize;
    169
    0
            let metadata_size = u32::from_le_bytes([data[20], data[21], data[22], data[23]]) as usize;
    170
    0
            let tensor_index_offset = u64::from_le_bytes([
    171
    0
                data[24], data[25], data[26], data[27], data[28], data[29], data[30], data[31],
    172
    0
            ]) as usize;
    173
    0
            let data_offset = u64::from_le_bytes([
    174
    0
                data[32], data[33], data[34], data[35], data[36], data[37], data[38], data[39],
    175
    0
            ]) as usize;
    176
    177
            // Parse metadata (JSON)
    178
    0
            let metadata_end = metadata_offset + metadata_size;
    179
    0
            if metadata_end > data.len() {
    180
    0
                return Err(RealizarError::FormatError {
    181
    0
                    reason: "Metadata extends beyond file".to_string(),
    182
    0
                });
    183
    0
            }
    184
    185
    0
            let metadata_json = &data[metadata_offset..metadata_end];
    186
    0
            let metadata: serde_json::Value = serde_json::from_slice(metadata_json).unwrap_or_default();
    187
    188
            // Extract architecture info from metadata
    189
    0
            let hidden_dim = metadata
    190
    0
                .get("hidden_size")
    191
    0
                .or_else(|| metadata.get("hidden_dim"))
    192
    0
                .and_then(serde_json::Value::as_u64)
    193
    0
                .unwrap_or(64) as usize;
    194
    195
    0
            let num_layers = metadata
    196
    0
                .get("num_hidden_layers")
    197
    0
                .or_else(|| metadata.get("num_layers"))
    198
    0
                .and_then(serde_json::Value::as_u64)
    199
    0
                .unwrap_or(1) as usize;
    200
    201
    0
            let num_heads = metadata
    202
    0
                .get("num_attention_heads")
    203
    0
                .or_else(|| metadata.get("num_heads"))
    204
    0
                .and_then(serde_json::Value::as_u64)
    205
    0
                .unwrap_or(4) as usize;
    206
    207
    0
            let num_kv_heads = metadata
    208
    0
                .get("num_key_value_heads")
    209
    0
                .or_else(|| metadata.get("num_kv_heads"))
    210
    0
                .and_then(serde_json::Value::as_u64)
    211
    0
                .unwrap_or(num_heads as u64) as usize;
    212
    213
    0
            let vocab_size = metadata
    214
    0
                .get("vocab_size")
    215
    0
                .and_then(serde_json::Value::as_u64)
    216
    0
                .unwrap_or(32000) as usize;
    217
    218
    0
            let intermediate_dim = metadata
    219
    0
                .get("intermediate_size")
    220
    0
                .or_else(|| metadata.get("intermediate_dim"))
    221
    0
                .and_then(serde_json::Value::as_u64)
    222
    0
                .unwrap_or((hidden_dim * 4) as u64) as usize;
    223
    224
    0
            let rope_theta = metadata
    225
    0
                .get("rope_theta")
    226
    0
                .and_then(serde_json::Value::as_f64)
    227
    0
                .unwrap_or(10000.0) as f32;
    228
    229
    0
            let rms_norm_eps = metadata
    230
    0
                .get("rms_norm_eps")
    231
    0
                .and_then(serde_json::Value::as_f64)
    232
    0
                .unwrap_or(1e-6) as f32;
    233
    234
    0
            let max_position = metadata
    235
    0
                .get("max_position_embeddings")
    236
    0
                .and_then(serde_json::Value::as_u64)
    237
    0
                .unwrap_or(2048) as usize;
    238
    239
    0
            let config = AprTransformerConfig {
    240
    0
                hidden_dim,
    241
    0
                num_layers,
    242
    0
                num_heads,
    243
    0
                num_kv_heads,
    244
    0
                vocab_size,
    245
    0
                intermediate_dim,
    246
    0
                context_length: max_position,
    247
    0
                rope_theta,
    248
    0
                eps: rms_norm_eps,
    249
    0
                ..Default::default()
    250
    0
            };
    251
    252
            // Parse tensor index
    253
            // APR v2 TensorIndexEntry format:
    254
            //   - name_len (2 bytes) + name (variable)
    255
            //   - dtype (1 byte)
    256
            //   - ndim (1 byte) + dims (8 bytes each)
    257
            //   - offset (8 bytes)
    258
            //   - size (8 bytes)
    259
            // Tuple: (offset, size, dims, dtype)
    260
    0
            let mut tensors: std::collections::BTreeMap<String, (usize, usize, Vec<usize>, u8)> =
    261
    0
                std::collections::BTreeMap::new();
    262
    263
    0
            let mut pos = tensor_index_offset;
    264
    0
            for _ in 0..tensor_count {
    265
    0
                if pos + 4 > data.len() {
    266
    0
                    break;
    267
    0
                }
    268
    269
                // Read tensor name length and name
    270
    0
                let name_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
    271
    0
                pos += 2;
    272
    273
    0
                if pos + name_len + 18 > data.len() {
    274
    0
                    break;
    275
    0
                }
    276
    277
    0
                let name = String::from_utf8_lossy(&data[pos..pos + name_len]).to_string();
    278
    0
                pos += name_len;
    279
    280
                // Read dtype (1 byte)
    281
    0
                let dtype = data[pos];
    282
    0
                pos += 1;
    283
    284
                // Read ndim (1 byte)
    285
    0
                let ndim = data[pos] as usize;
    286
    0
                pos += 1;
    287
    288
                // Read dimensions (8 bytes each)
    289
    0
                let mut dims = Vec::with_capacity(ndim);
    290
    0
                for _ in 0..ndim {
    291
    0
                    if pos + 8 > data.len() {
    292
    0
                        break;
    293
    0
                    }
    294
    0
                    let dim = u64::from_le_bytes([
    295
    0
                        data[pos],
    296
    0
                        data[pos + 1],
    297
    0
                        data[pos + 2],
    298
    0
                        data[pos + 3],
    299
    0
                        data[pos + 4],
    300
    0
                        data[pos + 5],
    301
    0
                        data[pos + 6],
    302
    0
                        data[pos + 7],
    303
    0
                    ]) as usize;
    304
    0
                    dims.push(dim);
    305
    0
                    pos += 8;
    306
                }
    307
    308
                // Read offset (8 bytes)
    309
    0
                if pos + 16 > data.len() {
    310
    0
                    break;
    311
    0
                }
    312
    0
                let offset = u64::from_le_bytes([
    313
    0
                    data[pos],
    314
    0
                    data[pos + 1],
    315
    0
                    data[pos + 2],
    316
    0
                    data[pos + 3],
    317
    0
                    data[pos + 4],
    318
    0
                    data[pos + 5],
    319
    0
                    data[pos + 6],
    320
    0
                    data[pos + 7],
    321
    0
                ]) as usize;
    322
    0
                pos += 8;
    323
    324
                // Read size (8 bytes)
    325
    0
                let size = u64::from_le_bytes([
    326
    0
                    data[pos],
    327
    0
                    data[pos + 1],
    328
    0
                    data[pos + 2],
    329
    0
                    data[pos + 3],
    330
    0
                    data[pos + 4],
    331
    0
                    data[pos + 5],
    332
    0
                    data[pos + 6],
    333
    0
                    data[pos + 7],
    334
    0
                ]) as usize;
    335
    0
                pos += 8;
    336
    337
    0
                tensors.insert(name, (data_offset + offset, size, dims, dtype));
    338
            }
    339
    340
            // Helper to extract f32 tensor (with Q4_K dequantization support)
    341
    0
            let get_f32_tensor = |name: &str| -> Option<Vec<f32>> {
    342
    0
                tensors.get(name).map(|(offset, size, dims, dtype)| {
    343
    0
                    let end = offset + size;
    344
    0
                    if end > data.len() {
    345
    0
                        return Vec::new();
    346
    0
                    }
    347
    0
                    let tensor_data = &data[*offset..end];
    348
    349
    0
                    match dtype {
    350
                        // Q4_K: converter dtype=8 or APR v2 native dtype=12
    351
                        8 | 12 => {
    352
    0
                            let num_elements: usize = dims.iter().product();
    353
    0
                            dequantize_q4_k_apr(tensor_data, num_elements)
    354
                        }
    355
                        // Q6_K: converter dtype=9 or APR v2 native dtype=14
    356
                        9 | 14 => {
    357
    0
                            let num_elements: usize = dims.iter().product();
    358
    0
                            dequantize_q6_k_apr(tensor_data, num_elements)
    359
                        }
    360
                        // F32 (dtype=0) or other: interpret as raw F32
    361
    0
                        _ => tensor_data
    362
    0
                            .chunks_exact(4)
    363
    0
                            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
    364
    0
                            .collect(),
    365
                    }
    366
    0
                })
    367
    0
            };
    368
    369
            // PMAT-103 FIX: Helper to get raw Q4K bytes (no dequantization) for fused kernel
    370
            // Returns None if tensor is not Q4K/Q5K/Q6K format
    371
            // APR dtype mapping (from GgufToAprQ4KConverter):
    372
            //   8 = Q4_K (GGML 12) or Q5_K (GGML 13)
    373
            //   9 = Q6_K (GGML 14)
    374
            //  10 = Q8_0 (GGML 8)
    375
            //  12 = Q4_K (APR v2 native)
    376
    0
            let get_q4k_raw_bytes = |name: &str| -> Option<Vec<u8>> {
    377
    0
                tensors.get(name).and_then(|(offset, size, _dims, dtype)| {
    378
                    // Accept Q4K tensors from either converter dtype (8) or APR v2 native dtype (12)
    379
    0
                    if *dtype != 8 && *dtype != 12 {
    380
    0
                        return None;
    381
    0
                    }
    382
    0
                    let end = offset + size;
    383
    0
                    if end > data.len() {
    384
    0
                        return None;
    385
    0
                    }
    386
    0
                    Some(data[*offset..end].to_vec())
    387
    0
                })
    388
    0
            };
    389
    390
            // PMAT-103 FIX: Also extract Q6K raw bytes for fused kernel
    391
    0
            let get_q6k_raw_bytes = |name: &str| -> Option<Vec<u8>> {
    392
    0
                tensors.get(name).and_then(|(offset, size, _dims, dtype)| {
    393
                    // Accept Q6K tensors from either converter dtype (9) or APR v2 native dtype (14)
    394
    0
                    if *dtype != 9 && *dtype != 14 {
    395
    0
                        return None;
    396
    0
                    }
    397
    0
                    let end = offset + size;
    398
    0
                    if end > data.len() {
    399
    0
                        return None;
    400
    0
                    }
    401
    0
                    Some(data[*offset..end].to_vec())
    402
    0
                })
    403
    0
            };
    404
    405
            // Debug: print available tensor names
    406
    0
            eprintln!("[DEBUG] APR v2 tensor count: {tensor_count}");
    407
    0
            eprintln!("[DEBUG] Available tensor names (first 10):");
    408
    0
            for (i, (name, (offset, size, dims, dtype))) in tensors.iter().enumerate() {
    409
    0
                if i < 10 {
    410
    0
                    eprintln!("  {name}: offset={offset}, size={size}, dims={dims:?}, dtype={dtype}");
    411
    0
                }
    412
            }
    413
    414
            // PMAT-086 FIX: Transpose matrix from GGUF [in_dim, out_dim] to matmul [out_dim, in_dim]
    415
            // GGUF/APR stores weights as [rows, cols] = [in_dim, out_dim] for y = x @ W
    416
            // But our matmul expects [out_dim, in_dim] for y = W @ x (row-major GEMV)
    417
    0
            let _transpose_weight = |data: Vec<f32>, rows: usize, cols: usize| -> Vec<f32> {
    418
    0
                let mut transposed = vec![0.0f32; rows * cols];
    419
    0
                for r in 0..rows {
    420
    0
                    for c in 0..cols {
    421
                        // data[r, c] -> transposed[c, r]
    422
    0
                        let src_idx = r * cols + c;
    423
    0
                        let dst_idx = c * rows + r;
    424
    0
                        if src_idx < data.len() && dst_idx < transposed.len() {
    425
    0
                            transposed[dst_idx] = data[src_idx];
    426
    0
                        }
    427
                    }
    428
                }
    429
    0
                transposed
    430
    0
            };
    431
    432
            // PMAT-086: Detect if using GGUF naming (output.weight, blk.X) or HF naming (lm_head.weight)
    433
            // GGUF uses [hidden_dim, vocab_size], HF uses [vocab_size, hidden_dim]
    434
    0
            let is_gguf_model = tensors.contains_key("output.weight") || tensors.contains_key("blk.0.attn_q.weight");
    435
    0
            eprintln!("[DEBUG] is_gguf_model={is_gguf_model}");
    436
    437
            // PMAT-086: Debug - check which embedding tensor names exist
    438
    0
            let embed_names = ["model.embed_tokens.weight", "token_embd.weight", "tok_embeddings.weight"];
    439
    0
            for name in &embed_names {
    440
    0
                if let Some((offset, size, dims, dtype)) = tensors.get(*name) {
    441
    0
                    eprintln!("[DEBUG] Found embedding {name}: offset={offset}, size={size}, dims={dims:?}, dtype={dtype}");
    442
    0
                }
    443
            }
    444
    445
            // Try to load token embedding
    446
    0
            let token_embedding_raw = get_f32_tensor("model.embed_tokens.weight")
    447
    0
                .or_else(|| get_f32_tensor("token_embd.weight"))
    448
    0
                .or_else(|| get_f32_tensor("tok_embeddings.weight"))
    449
    0
                .unwrap_or_else(|| {
    450
    0
                    eprintln!("[DEBUG] WARNING: No embedding tensor found! Using zeros.");
    451
    0
                    vec![0.0; vocab_size * hidden_dim]
    452
    0
                });
    453
    454
            // PMAT-086 FIX: APR stores GGUF data in row-major [vocab_size, hidden_dim] layout
    455
            // even though the dims metadata says [hidden_dim, vocab_size] (GGML column-major convention)
    456
            // The data is already correct - DO NOT transpose!
    457
    0
            let token_embedding = token_embedding_raw;
    458
    459
    0
            eprintln!("[DEBUG] token_embedding loaded: {} elements, first 5: {:?}",
    460
    0
                      token_embedding.len(),
    461
    0
                      &token_embedding[..5.min(token_embedding.len())]);
    462
    463
            // Load output norm
    464
    0
            let output_norm_weight = get_f32_tensor("model.norm.weight")
    465
    0
                .or_else(|| get_f32_tensor("output_norm.weight"))
    466
    0
                .unwrap_or_else(|| vec![1.0; hidden_dim]);
    467
    468
            // Debug: check output.weight / lm_head.weight
    469
    0
            for name in &["output.weight", "lm_head.weight"] {
    470
    0
                if let Some((offset, size, dims, dtype)) = tensors.get(*name) {
    471
    0
                    eprintln!("[DEBUG] Found lm_head {name}: offset={offset}, size={size}, dims={dims:?}, dtype={dtype}");
    472
    0
                }
    473
            }
    474
    475
            // Load LM head
    476
            // For tied embeddings (common in Qwen, LLaMA models), use embed_tokens as fallback
    477
    0
            let lm_head_raw = get_f32_tensor("lm_head.weight")
    478
    0
                .or_else(|| get_f32_tensor("output.weight"))
    479
    0
                .or_else(|| {
    480
                    // Weight tying: use embedding weights for lm_head
    481
    0
                    eprintln!("[DEBUG] Using tied weights: embedding -> lm_head");
    482
    0
                    get_f32_tensor("model.embed_tokens.weight")
    483
    0
                })
    484
    0
                .or_else(|| get_f32_tensor("token_embd.weight"))
    485
    0
                .unwrap_or_else(|| {
    486
    0
                    eprintln!("[DEBUG] WARNING: No lm_head tensor found! Using zeros.");
    487
    0
                    vec![0.0; hidden_dim * vocab_size]
    488
    0
                });
    489
    0
            eprintln!("[DEBUG] lm_head_raw: {} elements, first 5: {:?}",
    490
    0
                      lm_head_raw.len(),
    491
    0
                      &lm_head_raw[..5.min(lm_head_raw.len())]);
    492
            // PMAT-086 FIX: APR stores GGUF data in row-major [vocab_size, hidden_dim] layout
    493
            // even though the dims metadata says [hidden_dim, vocab_size] (GGML column-major convention)
    494
            // The data is already correct - DO NOT transpose!
    495
    0
            let lm_head_weight = lm_head_raw;
    496
    497
            // PMAT-103: Load lm_head Q4K/Q6K raw bytes for fused kernel inference
    498
    0
            let lm_head_weight_q4k = get_q4k_raw_bytes("lm_head.weight")
    499
    0
                .or_else(|| get_q4k_raw_bytes("output.weight"));
    500
    0
            let lm_head_weight_q6k = get_q6k_raw_bytes("lm_head.weight")
    501
    0
                .or_else(|| get_q6k_raw_bytes("output.weight"));
    502
    0
            if lm_head_weight_q4k.is_some() {
    503
    0
                eprintln!("[DEBUG] Loaded lm_head Q4K raw bytes for fused kernel");
    504
    0
            } else if lm_head_weight_q6k.is_some() {
    505
    0
                eprintln!("[DEBUG] Loaded lm_head Q6K raw bytes for fused kernel");
    506
    0
            }
    507
    508
            // Compute KV dimension from config
    509
    0
            let head_dim = hidden_dim / num_heads;
    510
    0
            let kv_dim = num_kv_heads * head_dim;
    511
    512
            // Load layers
    513
    0
            let mut layers = Vec::with_capacity(num_layers);
    514
            // PMAT-103 FIX: Also extract Q4K raw bytes for fused kernel inference
    515
    0
            let mut q4k_layer_weights: Vec<Q4KLayerWeights> = Vec::with_capacity(num_layers);
    516
    0
            let mut has_any_q4k = false;
    517
    518
    0
            for i in 0..num_layers {
    519
    0
                let hf_prefix = format!("model.layers.{i}");
    520
    0
                let gguf_prefix = format!("blk.{i}");
    521
    522
                // Try separate Q/K/V or combined QKV
    523
                // Support both HuggingFace and GGUF naming conventions
    524
                // PMAT-086 FIX: HF uses [out_dim, in_dim], GGUF uses [in_dim, out_dim]
    525
                // Only transpose GGUF tensors, not HF tensors
    526
    0
                let qkv_out_dim = hidden_dim + kv_dim + kv_dim;
    527
    528
                // Detect if using GGUF naming (blk.X) or HF naming (model.layers.X)
    529
    0
                let is_gguf = tensors.contains_key(&format!("{gguf_prefix}.attn_q.weight"));
    530
    531
    0
                let qkv_weight = if let Some(qkv) =
    532
    0
                    get_f32_tensor(&format!("{hf_prefix}.self_attn.qkv_proj.weight"))
    533
                {
    534
                    // HF fused QKV - already in [qkv_out_dim, hidden_dim] format
    535
    0
                    qkv
    536
                } else {
    537
                    // Get Q weight
    538
    0
                    let q_raw = get_f32_tensor(&format!("{hf_prefix}.self_attn.q_proj.weight"))
    539
    0
                        .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_q.weight")))
    540
    0
                        .unwrap_or_else(|| vec![0.0; hidden_dim * hidden_dim]);
    541
                    // Get K weight
    542
    0
                    let k_raw = get_f32_tensor(&format!("{hf_prefix}.self_attn.k_proj.weight"))
    543
    0
                        .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_k.weight")))
    544
    0
                        .unwrap_or_else(|| vec![0.0; hidden_dim * kv_dim]);
    545
                    // Get V weight
    546
    0
                    let v_raw = get_f32_tensor(&format!("{hf_prefix}.self_attn.v_proj.weight"))
    547
    0
                        .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_v.weight")))
    548
    0
                        .unwrap_or_else(|| vec![0.0; hidden_dim * kv_dim]);
    549
    550
                    // PMAT-086 FIX: Both HF and GGUF data are in [out_dim, in_dim] layout
    551
                    // GGUF dims say [in_dim, out_dim] but data is actually [out_dim, in_dim] due to GGML column-major convention
    552
                    // Fuse Q, K, V by stacking rows (Q, then K, then V) - no transpose needed
    553
    0
                    let _ = is_gguf; // Suppress unused warning
    554
    0
                    let mut qkv = Vec::with_capacity(qkv_out_dim * hidden_dim);
    555
    0
                    qkv.extend_from_slice(&q_raw);
    556
    0
                    qkv.extend_from_slice(&k_raw);
    557
    0
                    qkv.extend_from_slice(&v_raw);
    558
    0
                    qkv
    559
                };
    560
    561
                // Get Q/K/V biases (optional, for Qwen models)
    562
    0
                let q_bias = get_f32_tensor(&format!("{hf_prefix}.self_attn.q_proj.bias"))
    563
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_q.bias")));
    564
    0
                let k_bias = get_f32_tensor(&format!("{hf_prefix}.self_attn.k_proj.bias"))
    565
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_k.bias")));
    566
    0
                let v_bias = get_f32_tensor(&format!("{hf_prefix}.self_attn.v_proj.bias"))
    567
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_v.bias")));
    568
    569
                // Combine biases if present
    570
    0
                let qkv_bias = match (&q_bias, &k_bias, &v_bias) {
    571
    0
                    (Some(q), Some(k), Some(v)) => {
    572
    0
                        let mut bias = Vec::with_capacity(qkv_out_dim);
    573
    0
                        bias.extend_from_slice(q);
    574
    0
                        bias.extend_from_slice(k);
    575
    0
                        bias.extend_from_slice(v);
    576
    0
                        Some(bias)
    577
                    },
    578
    0
                    _ => None,
    579
                };
    580
    581
                // PMAT-086 FIX: Both HF and GGUF data are in [out_dim, in_dim] layout - no transpose needed
    582
    0
                let attn_output = get_f32_tensor(&format!("{hf_prefix}.self_attn.o_proj.weight"))
    583
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_output.weight")))
    584
    0
                    .unwrap_or_else(|| vec![0.0; hidden_dim * hidden_dim]);
    585
    586
    0
                let attn_norm = get_f32_tensor(&format!("{hf_prefix}.input_layernorm.weight"))
    587
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.attn_norm.weight")))
    588
    0
                    .unwrap_or_else(|| vec![1.0; hidden_dim]);
    589
    590
    0
                let ffn_norm = get_f32_tensor(&format!("{hf_prefix}.post_attention_layernorm.weight"))
    591
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.ffn_norm.weight")));
    592
    593
                // PMAT-086 FIX: FFN weights - both HF and GGUF data are in [out_dim, in_dim] layout
    594
                // No transpose needed - GGML column-major dims but row-major data
    595
    0
                let ffn_gate = get_f32_tensor(&format!("{hf_prefix}.mlp.gate_proj.weight"))
    596
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.ffn_gate.weight")));
    597
    0
                let ffn_up = get_f32_tensor(&format!("{hf_prefix}.mlp.up_proj.weight"))
    598
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.ffn_up.weight")))
    599
    0
                    .unwrap_or_else(|| vec![0.0; hidden_dim * intermediate_dim]);
    600
    0
                let ffn_down = get_f32_tensor(&format!("{hf_prefix}.mlp.down_proj.weight"))
    601
    0
                    .or_else(|| get_f32_tensor(&format!("{gguf_prefix}.ffn_down.weight")))
    602
    0
                    .unwrap_or_else(|| vec![0.0; intermediate_dim * hidden_dim]);
    603
    604
                // PMAT-103 FIX: Extract Q4K and Q6K raw bytes for fused kernel
    605
                // Now includes separate Q/K/V weights for fused QKV projection
    606
    0
                let q4k_attn_q = get_q4k_raw_bytes(&format!("{gguf_prefix}.attn_q.weight"));
    607
    0
                let q4k_attn_k = get_q4k_raw_bytes(&format!("{gguf_prefix}.attn_k.weight"));
    608
    0
                let q4k_attn_v = get_q4k_raw_bytes(&format!("{gguf_prefix}.attn_v.weight"));
    609
    0
                let q6k_attn_v = get_q6k_raw_bytes(&format!("{gguf_prefix}.attn_v.weight"));
    610
    0
                let q4k_attn_output = get_q4k_raw_bytes(&format!("{gguf_prefix}.attn_output.weight"));
    611
    0
                let q4k_ffn_gate = get_q4k_raw_bytes(&format!("{gguf_prefix}.ffn_gate.weight"));
    612
    0
                let q4k_ffn_up = get_q4k_raw_bytes(&format!("{gguf_prefix}.ffn_up.weight"));
    613
    0
                let q4k_ffn_down = get_q4k_raw_bytes(&format!("{gguf_prefix}.ffn_down.weight"));
    614
                // Q6K fallback for tensors that aren't Q4K (common in mixed quantization models)
    615
    0
                let q6k_ffn_down = get_q6k_raw_bytes(&format!("{gguf_prefix}.ffn_down.weight"));
    616
    0
                let q6k_ffn_up = get_q6k_raw_bytes(&format!("{gguf_prefix}.ffn_up.weight"));
    617
    618
    0
                let has_q4k_weights = q4k_attn_q.is_some()
    619
    0
                    || q4k_attn_k.is_some()
    620
    0
                    || q4k_attn_output.is_some()
    621
    0
                    || q4k_ffn_gate.is_some()
    622
    0
                    || q4k_ffn_up.is_some()
    623
    0
                    || q4k_ffn_down.is_some();
    624
    0
                let has_q6k_weights = q6k_ffn_down.is_some() || q6k_ffn_up.is_some() || q6k_attn_v.is_some();
    625
    626
    0
                if has_q4k_weights || has_q6k_weights {
    627
    0
                    has_any_q4k = true;
    628
    0
                }
    629
    630
    0
                q4k_layer_weights.push(Q4KLayerWeights {
    631
    0
                    qkv_weight: None, // Q+K+V are separate tensors, not combined
    632
    0
                    attn_q_weight: q4k_attn_q,
    633
    0
                    attn_k_weight: q4k_attn_k,
    634
    0
                    attn_v_weight: q4k_attn_v,
    635
    0
                    attn_v_weight_q6k: q6k_attn_v,
    636
    0
                    attn_output_weight: q4k_attn_output,
    637
    0
                    ffn_gate_weight: q4k_ffn_gate,
    638
    0
                    ffn_up_weight: q4k_ffn_up,
    639
    0
                    ffn_down_weight: q4k_ffn_down,
    640
    0
                    ffn_down_weight_q6k: q6k_ffn_down,
    641
    0
                    ffn_up_weight_q6k: q6k_ffn_up,
    642
    0
                });
    643
    644
    0
                layers.push(AprTransformerLayer {
    645
    0
                    attn_norm_weight: attn_norm,
    646
    0
                    attn_norm_bias: None,
    647
    0
                    qkv_weight,
    648
    0
                    qkv_bias,
    649
    0
                    attn_output_weight: attn_output,
    650
    0
                    attn_output_bias: None,
    651
    0
                    ffn_gate_weight: ffn_gate,
    652
    0
                    ffn_gate_bias: None,
    653
    0
                    ffn_up_weight: ffn_up,
    654
    0
                    ffn_up_bias: None,
    655
    0
                    ffn_down_weight: ffn_down,
    656
    0
                    ffn_down_bias: None,
    657
    0
                    ffn_norm_weight: ffn_norm,
    658
    0
                    ffn_norm_bias: None,
    659
    0
                });
    660
            }
    661
    662
            // PMAT-103 FIX: Store Q4K layer weights for fused kernel inference
    663
    0
            let q4k_layers = if has_any_q4k {
    664
    0
                eprintln!("[DEBUG] Loaded Q4K raw bytes for fused kernel inference");
    665
    0
                Some(q4k_layer_weights)
    666
            } else {
    667
    0
                None
    668
            };
    669
    670
    0
            Ok(Self {
    671
    0
                config,
    672
    0
                token_embedding,
    673
    0
                layers,
    674
    0
                output_norm_weight,
    675
    0
                output_norm_bias: None,
    676
    0
                lm_head_weight,
    677
    0
                lm_head_bias: None,
    678
    0
                q4k_layers,
    679
    0
                lm_head_weight_q6k,
    680
    0
                lm_head_weight_q4k,
    681
    0
            })
    682
    0
        }
    683
    684
        /// Create a new APR transformer with the given configuration
    685
    0
        pub fn new(config: AprTransformerConfig) -> Self {
    686
    0
            let hidden_dim = config.hidden_dim;
    687
    0
            let vocab_size = config.vocab_size;
    688
    0
            let intermediate_dim = config.intermediate_dim;
    689
    690
    0
            let layers = (0..config.num_layers)
    691
    0
                .map(|_| AprTransformerLayer::empty(hidden_dim, intermediate_dim))
    692
    0
                .collect();
    693
    694
    0
            Self {
    695
    0
                config,
    696
    0
                token_embedding: vec![0.0; vocab_size * hidden_dim],
    697
    0
                layers,
    698
    0
                output_norm_weight: vec![1.0; hidden_dim],
    699
    0
                output_norm_bias: None,
    700
    0
                lm_head_weight: vec![0.0; hidden_dim * vocab_size],
    701
    0
                lm_head_bias: None,
    702
    0
                q4k_layers: None,
    703
    0
                lm_head_weight_q6k: None,
    704
    0
                lm_head_weight_q4k: None,
    705
    0
            }
    706
    0
        }
    707
    708
        /// Get the model configuration
    709
        #[must_use]
    710
    0
        pub fn config(&self) -> &AprTransformerConfig {
    711
    0
            &self.config
    712
    0
        }
    713
    714
        /// Generate tokens autoregressively (simplified version without KV cache)
    715
        ///
    716
        /// # Arguments
    717
        ///
    718
        /// * `prompt` - Initial token IDs
    719
        /// * `max_tokens` - Maximum tokens to generate
    720
        ///
    721
        /// # Returns
    722
        ///
    723
        /// Generated token sequence (including prompt)
    724
    0
        pub fn generate(&self, prompt: &[u32], max_tokens: usize) -> Result<Vec<u32>> {
    725
    0
            let mut tokens = prompt.to_vec();
    726
    727
    0
            for _ in 0..max_tokens {
    728
    0
                let logits = self.forward(&tokens)?;
    729
    730
                // Greedy sampling: take argmax
    731
    0
                let next_token = logits
    732
    0
                    .iter()
    733
    0
                    .enumerate()
    734
    0
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    735
    0
                    .map_or(0, |(idx, _)| idx as u32);
    736
    737
    0
                tokens.push(next_token);
    738
    739
                // Stop at EOS tokens:
    740
                // - Standard: 2
    741
                // - Qwen2: 151645 (EOS), 151643 (BOS)
    742
                // - LLaMA: 2
    743
    0
                if next_token == 2 || next_token == 151645 || next_token == 151643 {
    744
    0
                    break;
    745
    0
                }
    746
            }
    747
    748
    0
            Ok(tokens)
    749
    0
        }
    750
    751
        /// Get total number of parameters
    752
        #[must_use]
    753
    12
        pub fn num_parameters(&self) -> usize {
    754
    12
            let mut count = 0;
    755
    12
            count += self.token_embedding.len();
    756
    34
            for layer22 in &self.layers {
    757
    22
                count += layer.num_parameters();
    758
    22
            }
    759
    12
            count += self.output_norm_weight.len();
    760
    12
            count += self.output_norm_bias.as_ref().map_or(0, Vec::len);
    761
    12
            count += self.lm_head_weight.len();
    762
    12
            count += self.lm_head_bias.as_ref().map_or(0, Vec::len);
    763
    12
            count
    764
    12
        }
    765
    766
        /// Get memory size in bytes (F32 = 4 bytes per param)
    767
        #[must_use]
    768
    6
        pub fn memory_size(&self) -> usize {
    769
    6
            self.num_parameters() * 4
    770
    6
        }
    771
    772
        /// Look up token embeddings
    773
        #[must_use]
    774
    3
        pub fn embed(&self, token_ids: &[u32]) -> Vec<f32> {
    775
    3
            let hidden_dim = self.config.hidden_dim;
    776
    3
            let mut embeddings = Vec::with_capacity(token_ids.len() * hidden_dim);
    777
    778
    12
            for &token_id9 in token_ids {
    779
    9
                let offset = (token_id as usize) * hidden_dim;
    780
    9
                if offset + hidden_dim <= self.token_embedding.len() {
    781
    9
                    embeddings.extend_from_slice(&self.token_embedding[offset..offset + hidden_dim]);
    782
    9
                } else {
    783
    0
                    // Out of vocab - return zeros
    784
    0
                    embeddings.extend(std::iter::repeat_n(0.0, hidden_dim));
    785
    0
                }
    786
            }
    787
    788
    3
            embeddings
    789
    3
        }
    790
    791
        /// RMSNorm (Root Mean Square Layer Normalization)
    792
        ///
    793
        /// Used by Qwen2, LLaMA, Mistral, and most modern LLMs.
    794
        /// Formula: output = x / sqrt(mean(x^2) + eps) * weight
    795
        ///
    796
        /// PMAT-094: Fixed five-whys root cause - was using LayerNorm (mean subtraction)
    797
        /// instead of RMSNorm which caused garbage output for Qwen2 models.
    798
    6
        fn layer_norm(
    799
    6
            &self,
    800
    6
            input: &[f32],
    801
    6
            weight: &[f32],
    802
    6
            bias: Option<&[f32]>,
    803
    6
            eps: f32,
    804
    6
        ) -> Vec<f32> {
    805
    6
            let hidden_dim = self.config.hidden_dim;
    806
    6
            let seq_len = input.len() / hidden_dim;
    807
    6
            let mut output = Vec::with_capacity(input.len());
    808
    809
    18
            for s in 0..seq_len6 {
    810
    18
                let start = s * hidden_dim;
    811
    18
                let slice = &input[start..start + hidden_dim];
    812
    813
                // RMSNorm: compute root mean square (no mean subtraction!)
    814
    72
                let sum_sq18: f3218 = slice18.iter18().map18(|x| x * x).sum18();
    815
    18
                let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();
    816
    817
                // Normalize and scale
    818
    72
                for (i, &x) in slice18.iter18().enumerate18() {
    819
    72
                    let normalized = x / rms;
    820
    72
                    let scaled = normalized * weight[i];
    821
    72
                    let shifted = if let Some(b0) = bias {
    822
    0
                        scaled + b[i]
    823
                    } else {
    824
    72
                        scaled
    825
                    };
    826
    72
                    output.push(shifted);
    827
                }
    828
            }
    829
    830
    6
            output
    831
    6
        }
    832
    833
        /// Matrix multiplication: output[out_dim] = weight[out_dim, in_dim] @ input[in_dim]
    834
        ///
    835
        /// PMAT-095 FIX: Weights are now stored in matvec-optimal [out_dim, in_dim] format.
    836
        ///
    837
        /// PMAT-103 FIX: Zero-copy implementation using raw slice operations.
    838
        /// Previous implementations had O(n) allocation overhead per matmul call.
    839
        #[allow(clippy::unused_self)]
    840
    15
        fn matmul(&self, input: &[f32], weight: &[f32], in_dim: usize, out_dim: usize) -> Vec<f32> {
    841
    15
            let seq_len = input.len() / in_dim;
    842
    15
            let expected_size = in_dim * out_dim;
    843
    844
    15
            if weight.len() != expected_size {
    845
    0
                return self.matmul_scalar(input, weight, in_dim, out_dim);
    846
    15
            }
    847
    848
    15
            let mut output = vec![0.0f32; seq_len * out_dim];
    849
    850
    39
            for s in 0..seq_len15 {
    851
    39
                let input_start = s * in_dim;
    852
    39
                let input_slice = &input[input_start..input_start + in_dim];
    853
    39
                let out_start = s * out_dim;
    854
    855
                // PMAT-103: Unrolled dot product for better cache utilization
    856
                // Process 4 output elements at a time when possible
    857
    39
                let out_chunks = out_dim / 4;
    858
    39
                let out_remainder = out_dim % 4;
    859
    860
    69
                for o_chunk in 0..out_chunks39 {
    861
    69
                    let o_base = o_chunk * 4;
    862
    69
                    let mut sum0 = 0.0f32;
    863
    69
                    let mut sum1 = 0.0f32;
    864
    69
                    let mut sum2 = 0.0f32;
    865
    69
                    let mut sum3 = 0.0f32;
    866
    867
    69
                    let w0_start = (o_base) * in_dim;
    868
    69
                    let w1_start = (o_base + 1) * in_dim;
    869
    69
                    let w2_start = (o_base + 2) * in_dim;
    870
    69
                    let w3_start = (o_base + 3) * in_dim;
    871
    872
    312
                    for i in 0..in_dim69 {
    873
    312
                        let x = input_slice[i];
    874
    312
                        sum0 += x * weight[w0_start + i];
    875
    312
                        sum1 += x * weight[w1_start + i];
    876
    312
                        sum2 += x * weight[w2_start + i];
    877
    312
                        sum3 += x * weight[w3_start + i];
    878
    312
                    }
    879
    880
    69
                    output[out_start + o_base] = sum0;
    881
    69
                    output[out_start + o_base + 1] = sum1;
    882
    69
                    output[out_start + o_base + 2] = sum2;
    883
    69
                    output[out_start + o_base + 3] = sum3;
    884
                }
    885
    886
                // Handle remainder
    887
    39
                for o6 in (out_dim - out_remainder)..out_dim {
    888
    6
                    let w_start = o * in_dim;
    889
    6
                    let mut sum = 0.0f32;
    890
    24
                    for i in 0..in_dim6 {
    891
    24
                        sum += input_slice[i] * weight[w_start + i];
    892
    24
                    }
    893
    6
                    output[out_start + o] = sum;
    894
                }
    895
            }
    896
    897
    15
            output
    898
    15
        }
    899
    900
        /// Scalar fallback for matmul (used when trueno fails)
    901
        ///
    902
        /// PMAT-095: Weight is [out_dim, in_dim] row-major format
    903
        #[allow(clippy::unused_self)]
    904
    0
        fn matmul_scalar(
    905
    0
            &self,
    906
    0
            input: &[f32],
    907
    0
            weight: &[f32],
    908
    0
            in_dim: usize,
    909
    0
            out_dim: usize,
    910
    0
        ) -> Vec<f32> {
    911
    0
            let seq_len = input.len() / in_dim;
    912
    0
            let mut output = Vec::with_capacity(seq_len * out_dim);
    913
    914
    0
            for s in 0..seq_len {
    915
    0
                let input_start = s * in_dim;
    916
    0
                let input_slice = &input[input_start..input_start + in_dim];
    917
    918
    0
                for o in 0..out_dim {
    919
    0
                    let mut sum = 0.0;
    920
    0
                    for (i, &input_val) in input_slice.iter().enumerate() {
    921
                        // PMAT-095: Weight is [out_dim, in_dim] row-major
    922
    0
                        let weight_idx = o * in_dim + i;
    923
    0
                        if weight_idx < weight.len() {
    924
    0
                            sum += input_val * weight[weight_idx];
    925
    0
                        }
    926
                    }
    927
    0
                    output.push(sum);
    928
                }
    929
            }
    930
    931
    0
            output
    932
    0
        }
    933
    934
        /// Add bias in-place
    935
        #[allow(clippy::unused_self)]
    936
    0
        fn add_bias(&self, data: &mut [f32], bias: &[f32]) {
    937
    0
            let dim = bias.len();
    938
    0
            for (i, val) in data.iter_mut().enumerate() {
    939
    0
                *val += bias[i % dim];
    940
    0
            }
    941
    0
        }
    942
    943
        /// GELU activation (tanh approximation)
    944
        #[allow(clippy::unused_self)]
    945
    3
        fn gelu(&self, data: &mut [f32]) {
    946
            const SQRT_2_OVER_PI: f32 = 0.797_884_6;
    947
            const GELU_COEFF: f32 = 0.044_715;
    948
    949
    72
            for x in data3.iter_mut3() {
    950
    72
                let x3 = *x * *x * *x;
    951
    72
                let inner = SQRT_2_OVER_PI * (*x + GELU_COEFF * x3);
    952
    72
                *x = 0.5 * *x * (1.0 + inner.tanh());
    953
    72
            }
    954
    3
        }
    955
    956
        /// Apply Rotary Position Embedding (RoPE) to Q or K vectors
    957
        ///
    958
        /// RoPE encodes position information by rotating pairs of elements
    959
        /// with position-dependent angles.
    960
    18
        fn apply_rope_f32(&self, x: &mut [f32], position: usize, num_heads: usize, head_dim: usize) {
    961
    18
            let half_dim = head_dim / 2;
    962
    18
            let theta = self.config.rope_theta;
    963
    18
            let pos_f32 = position as f32;
    964
    18
            let head_dim_f32 = head_dim as f32;
    965
    966
    72
            for h in 0..num_heads18 {
    967
    72
                let head_start = h * head_dim;
    968
    72
                let idx2_start = head_start + half_dim;
    969
    970
    72
                if idx2_start + half_dim > x.len() {
    971
    0
                    continue;
    972
    72
                }
    973
    974
    72
                for i0 in 0..half_dim {
    975
    0
                    let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim_f32);
    976
    0
                    let angle = pos_f32 * freq;
    977
    0
                    let (sin_val, cos_val) = angle.sin_cos();
    978
    0
    979
    0
                    let x1 = x[head_start + i];
    980
    0
                    let x2 = x[idx2_start + i];
    981
    0
    982
    0
                    // Apply rotation: [cos -sin; sin cos] * [x1; x2]
    983
    0
                    x[head_start + i] = x1 * cos_val - x2 * sin_val;
    984
    0
                    x[idx2_start + i] = x1 * sin_val + x2 * cos_val;
    985
    0
                }
    986
            }
    987
    18
        }
    988
    989
        /// Forward pass through the transformer
    990
        ///
    991
        /// # Arguments
    992
        ///
    993
        /// * `token_ids` - Input token IDs
    994
        ///
    995
        /// # Returns
    996
        ///
    997
        /// Logits over vocabulary for next token prediction
    998
        ///
    999
        /// # Errors
    1000
        ///
    1001
        /// Returns error if inference fails
    1002
    3
        pub fn forward(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    1003
    3
            if token_ids.is_empty() {
    1004
    0
                return Err(RealizarError::InvalidShape {
    1005
    0
                    reason: "Token sequence cannot be empty".to_string(),
    1006
    0
                });
    1007
    3
            }
    1008
    1009
    3
            let hidden_dim = self.config.hidden_dim;
    1010
    3
            let intermediate_dim = self.config.intermediate_dim;
    1011
    1012
            // 1. Token embedding lookup
    1013
    3
            let mut hidden = self.embed(token_ids);
    1014
    1015
            // 2. Process through transformer layers
    1016
    3
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    1017
                // PMAT-103: Get Q4K weights for this layer (if available)
    1018
    3
                let q4k_layer = self.q4k_layers.as_ref().and_then(|l| l0.get0(layer_idx0));
    1019
    1020
                // 2a. Attention layer norm
    1021
    3
                let normed = self.layer_norm(
    1022
    3
                    &hidden,
    1023
    3
                    &layer.attn_norm_weight,
    1024
    3
                    layer.attn_norm_bias.as_deref(),
    1025
    3
                    self.config.eps,
    1026
                );
    1027
    1028
                // 2b. QKV projection
    1029
                // Calculate qkv_dim from actual weight size (handles GQA models)
    1030
    3
                let qkv_dim = layer.qkv_weight.len() / hidden_dim;
    1031
    3
                let mut qkv = self.matmul(&normed, &layer.qkv_weight, hidden_dim, qkv_dim);
    1032
    3
                if let Some(ref bias0) = layer.qkv_bias {
    1033
    0
                    self.add_bias(&mut qkv, bias);
    1034
    3
                }
    1035
    1036
                // 2c. Proper attention with GQA support and RoPE
    1037
    3
                let seq_len = token_ids.len();
    1038
    3
                let head_dim = hidden_dim / self.config.num_heads;
    1039
    3
                let num_kv_heads = self.config.num_kv_heads;
    1040
    3
                let kv_dim = num_kv_heads * head_dim;
    1041
    3
                let group_size = self.config.num_heads / num_kv_heads;
    1042
    3
                let scale = 1.0 / (head_dim as f32).sqrt();
    1043
    1044
                // Split QKV and apply RoPE
    1045
    3
                let mut q_all = Vec::with_capacity(seq_len * hidden_dim);
    1046
    3
                let mut k_all = Vec::with_capacity(seq_len * kv_dim);
    1047
    3
                let mut v_all = Vec::with_capacity(seq_len * kv_dim);
    1048
    1049
    9
                for s in 0..seq_len3 {
    1050
    9
                    let qkv_start = s * qkv_dim;
    1051
    9
    1052
    9
                    // Extract Q, K, V (layout: [Q..., K..., V...])
    1053
    9
                    let mut q_pos = qkv[qkv_start..qkv_start + hidden_dim].to_vec();
    1054
    9
                    let mut k_pos =
    1055
    9
                        qkv[qkv_start + hidden_dim..qkv_start + hidden_dim + kv_dim].to_vec();
    1056
    9
                    let v_pos =
    1057
    9
                        &qkv[qkv_start + hidden_dim + kv_dim..qkv_start + hidden_dim + 2 * kv_dim];
    1058
    9
    1059
    9
                    // Apply RoPE to Q and K
    1060
    9
                    self.apply_rope_f32(&mut q_pos, s, self.config.num_heads, head_dim);
    1061
    9
                    self.apply_rope_f32(&mut k_pos, s, num_kv_heads, head_dim);
    1062
    9
    1063
    9
                    q_all.extend_from_slice(&q_pos);
    1064
    9
                    k_all.extend_from_slice(&k_pos);
    1065
    9
                    v_all.extend_from_slice(v_pos);
    1066
    9
                }
    1067
    1068
                // Compute scaled dot-product attention with causal mask
    1069
    3
                let mut attn_out = vec![0.0f32; seq_len * hidden_dim];
    1070
    12
                for head in 0..self.config.num_heads3 {
    1071
    12
                    let kv_head = head / group_size;
    1072
    12
                    let q_head_offset = head * head_dim;
    1073
    12
                    let kv_head_offset = kv_head * head_dim;
    1074
    1075
    36
                    for i in 0..seq_len12 {
    1076
                        // Compute attention scores for this position
    1077
    36
                        let mut scores = Vec::with_capacity(i + 1);
    1078
    36
                        let q_start = i * hidden_dim + q_head_offset;
    1079
    1080
    72
                        for j in 0..=i36 {
    1081
                            // Only attend to positions <= current (causal mask)
    1082
    72
                            let k_start = j * kv_dim + kv_head_offset;
    1083
    72
                            let mut score = 0.0f32;
    1084
    72
                            for d in 0..head_dim {
    1085
    72
                                score += q_all[q_start + d] * k_all[k_start + d];
    1086
    72
                            }
    1087
    72
                            scores.push(score * scale);
    1088
                        }
    1089
    1090
                        // Softmax
    1091
    36
                        let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1092
    36
                        let mut exp_sum = 0.0f32;
    1093
    108
                        for s72 in &mut scores {
    1094
    72
                            *s = (*s - max_score).exp();
    1095
    72
                            exp_sum += *s;
    1096
    72
                        }
    1097
    36
                        if exp_sum > 0.0 {
    1098
    108
                            for s72 in &mut scores {
    1099
    72
                                *s /= exp_sum;
    1100
    72
                            }
    1101
    0
                        }
    1102
    1103
                        // Weighted sum of V
    1104
    36
                        let out_start = i * hidden_dim + q_head_offset;
    1105
    72
                        for (j, &weight) in scores.iter()36.enumerate36() {
    1106
    72
                            let v_start = j * kv_dim + kv_head_offset;
    1107
    72
                            for d in 0..head_dim {
    1108
    72
                                attn_out[out_start + d] += weight * v_all[v_start + d];
    1109
    72
                            }
    1110
                        }
    1111
                    }
    1112
                }
    1113
    1114
                // 2d. Attention output projection
    1115
                // PMAT-103: Use Q4K fused kernel when available
    1116
    3
                let mut attn_output = if let Some(q4k_bytes0) = q4k_layer
    1117
    3
                    .and_then(|q| q.attn_output_weight0.as_ref0())
    1118
                {
    1119
    0
                    if layer_idx == 0 {
    1120
    0
                        eprintln!("[TRACE] Layer {layer_idx}: attn_output using Q4K fused kernel");
    1121
    0
                    }
    1122
                    // Fused Q4K matmul: process each position separately
    1123
                    // PMAT-103: Use column-major kernel for GGUF layout
    1124
    0
                    let seq_len = token_ids.len();
    1125
    0
                    let mut output = Vec::with_capacity(seq_len * hidden_dim);
    1126
    0
                    for s in 0..seq_len {
    1127
    0
                        let input_slice = &attn_out[s * hidden_dim..(s + 1) * hidden_dim];
    1128
    0
                        let pos_out = matmul_q4k_rowmajor(q4k_bytes, input_slice, hidden_dim, hidden_dim);
    1129
    0
                        output.extend(pos_out);
    1130
    0
                    }
    1131
    0
                    output
    1132
                } else {
    1133
    3
                    if layer_idx == 0 {
    1134
    3
                        eprintln!("[TRACE] Layer {layer_idx}: attn_output using F32 fallback (slow!)");
    1135
    3
                    }0
    1136
    3
                    self.matmul(&attn_out, &layer.attn_output_weight, hidden_dim, hidden_dim)
    1137
                };
    1138
    3
                if let Some(ref bias0) = layer.attn_output_bias {
    1139
    0
                    self.add_bias(&mut attn_output, bias);
    1140
    3
                }
    1141
    1142
                // 2e. Residual connection
    1143
    36
                for i in 0..hidden3.len3() {
    1144
    36
                    hidden[i] += attn_output[i];
    1145
    36
                }
    1146
    1147
                // 2f. Apply FFN norm if present (post_attention_layernorm)
    1148
    3
                let ffn_input = if let Some(ref ffn_norm0) = layer.ffn_norm_weight {
    1149
    0
                    self.layer_norm(
    1150
    0
                        &hidden,
    1151
    0
                        ffn_norm,
    1152
    0
                        layer.ffn_norm_bias.as_deref(),
    1153
    0
                        self.config.eps,
    1154
                    )
    1155
                } else {
    1156
    3
                    hidden.clone()
    1157
                };
    1158
    1159
                // 2g. FFN projection (SwiGLU or standard GELU)
    1160
                // PMAT-103: Use Q4K fused kernel when available for FFN
    1161
    3
                let seq_len = token_ids.len();
    1162
    3
                let ffn_output = if let Some(ref _gate_weight0) = layer.ffn_gate_weight {
    1163
                    // SwiGLU: down(SiLU(gate(x)) * up(x))
    1164
                    // PMAT-103: Check for Q4K gate weight
    1165
    0
                    let gate = if let Some(q4k_bytes) = q4k_layer
    1166
    0
                        .and_then(|q| q.ffn_gate_weight.as_ref())
    1167
                    {
    1168
    0
                        if layer_idx == 0 {
    1169
    0
                            eprintln!("[TRACE] Layer {layer_idx}: ffn_gate using Q4K fused kernel");
    1170
    0
                        }
    1171
    0
                        let mut output = Vec::with_capacity(seq_len * intermediate_dim);
    1172
    0
                        for s in 0..seq_len {
    1173
    0
                            let input_slice = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    1174
    0
                            // PMAT-103 FIX: Q4K kernel expects (ne0=output_dim, ne1=input_dim)
    1175
    0
                            // ffn_gate: [intermediate_dim, hidden_dim] maps hidden[1536] -> intermediate[8960]
    1176
    0
                            let pos_out = matmul_q4k_rowmajor(q4k_bytes, input_slice, intermediate_dim, hidden_dim);
    1177
    0
                            output.extend(pos_out);
    1178
    0
                        }
    1179
    0
                        output
    1180
                    } else {
    1181
    0
                        self.matmul(&ffn_input, layer.ffn_gate_weight.as_ref().expect("gate weight"), hidden_dim, intermediate_dim)
    1182
                    };
    1183
    1184
                    // PMAT-103: Check for Q4K up weight
    1185
    0
                    let up = if let Some(q4k_bytes) = q4k_layer
    1186
    0
                        .and_then(|q| q.ffn_up_weight.as_ref())
    1187
                    {
    1188
    0
                        if layer_idx == 0 {
    1189
    0
                            eprintln!("[TRACE] Layer {layer_idx}: ffn_up using Q4K fused kernel");
    1190
    0
                        }
    1191
    0
                        let mut output = Vec::with_capacity(seq_len * intermediate_dim);
    1192
    0
                        for s in 0..seq_len {
    1193
    0
                            let input_slice = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    1194
    0
                            // PMAT-103 FIX: Q4K kernel expects (ne0=output_dim, ne1=input_dim)
    1195
    0
                            // ffn_up: [intermediate_dim, hidden_dim] maps hidden[1536] -> intermediate[8960]
    1196
    0
                            let pos_out = matmul_q4k_rowmajor(q4k_bytes, input_slice, intermediate_dim, hidden_dim);
    1197
    0
                            output.extend(pos_out);
    1198
    0
                        }
    1199
    0
                        output
    1200
                    } else {
    1201
    0
                        if layer_idx == 0 {
    1202
    0
                            eprintln!("[TRACE] Layer {layer_idx}: ffn_up using F32 fallback (slow!)");
    1203
    0
                        }
    1204
    0
                        self.matmul(&ffn_input, &layer.ffn_up_weight, hidden_dim, intermediate_dim)
    1205
                    };
    1206
    1207
                    // SiLU(gate) * up, then down projection
    1208
    0
                    let mut ffn_hidden = Vec::with_capacity(gate.len());
    1209
    0
                    for (g, u) in gate.iter().zip(up.iter()) {
    1210
    0
                        let silu_g = g / (1.0 + (-g).exp()); // SiLU = x * sigmoid(x)
    1211
    0
                        ffn_hidden.push(silu_g * u);
    1212
    0
                    }
    1213
    1214
                    // PMAT-103: Check for Q4K or Q6K down weight
    1215
    0
                    let mut out = if let Some(q4k_bytes) = q4k_layer
    1216
    0
                        .and_then(|q| q.ffn_down_weight.as_ref())
    1217
                    {
    1218
    0
                        if layer_idx == 0 {
    1219
    0
                            eprintln!("[TRACE] Layer {layer_idx}: ffn_down using Q4K fused kernel");
    1220
    0
                        }
    1221
    0
                        let mut output = Vec::with_capacity(seq_len * hidden_dim);
    1222
    0
                        for s in 0..seq_len {
    1223
    0
                            let input_slice = &ffn_hidden[s * intermediate_dim..(s + 1) * intermediate_dim];
    1224
    0
                            let pos_out = matmul_q4k_rowmajor(q4k_bytes, input_slice, hidden_dim, intermediate_dim);
    1225
    0
                            output.extend(pos_out);
    1226
    0
                        }
    1227
    0
                        output
    1228
    0
                    } else if let Some(q6k_bytes) = q4k_layer
    1229
    0
                        .and_then(|q| q.ffn_down_weight_q6k.as_ref())
    1230
                    {
    1231
    0
                        if layer_idx == 0 {
    1232
    0
                            eprintln!("[TRACE] Layer {layer_idx}: ffn_down using Q6K fused kernel");
    1233
    0
                        }
    1234
    0
                        let mut output = Vec::with_capacity(seq_len * hidden_dim);
    1235
    0
                        for s in 0..seq_len {
    1236
    0
                            let input_slice = &ffn_hidden[s * intermediate_dim..(s + 1) * intermediate_dim];
    1237
    0
                            let pos_out = matmul_q6k_rowmajor(q6k_bytes, input_slice, hidden_dim, intermediate_dim);
    1238
    0
                            output.extend(pos_out);
    1239
    0
                        }
    1240
    0
                        output
    1241
                    } else {
    1242
    0
                        if layer_idx == 0 {
    1243
    0
                            eprintln!("[TRACE] Layer {layer_idx}: ffn_down using F32 fallback (slow!)");
    1244
    0
                        }
    1245
    0
                        self.matmul(&ffn_hidden, &layer.ffn_down_weight, intermediate_dim, hidden_dim)
    1246
                    };
    1247
    0
                    if let Some(ref bias) = layer.ffn_down_bias {
    1248
    0
                        self.add_bias(&mut out, bias);
    1249
    0
                    }
    1250
    0
                    out
    1251
                } else {
    1252
                    // Standard MLP: down(GELU(up(x)))
    1253
                    // PMAT-103: Check for Q4K up weight
    1254
    3
                    let mut ffn_hidden = if let Some(q4k_bytes0) = q4k_layer
    1255
    3
                        .and_then(|q| q.ffn_up_weight0.as_ref0())
    1256
                    {
    1257
    0
                        let mut output = Vec::with_capacity(seq_len * intermediate_dim);
    1258
    0
                        for s in 0..seq_len {
    1259
    0
                            let input_slice = &ffn_input[s * hidden_dim..(s + 1) * hidden_dim];
    1260
    0
                            // PMAT-103 FIX: Q4K kernel expects (ne0=output_dim, ne1=input_dim)
    1261
    0
                            // ffn_up: [intermediate_dim, hidden_dim] maps hidden[1536] -> intermediate[8960]
    1262
    0
                            let pos_out = matmul_q4k_rowmajor(q4k_bytes, input_slice, intermediate_dim, hidden_dim);
    1263
    0
                            output.extend(pos_out);
    1264
    0
                        }
    1265
    0
                        output
    1266
                    } else {
    1267
    3
                        self.matmul(&ffn_input, &layer.ffn_up_weight, hidden_dim, intermediate_dim)
    1268
                    };
    1269
    3
                    if let Some(ref bias0) = layer.ffn_up_bias {
    1270
    0
                        self.add_bias(&mut ffn_hidden, bias);
    1271
    3
                    }
    1272
    3
                    self.gelu(&mut ffn_hidden);
    1273
    1274
                    // PMAT-103: Check for Q4K down weight
    1275
    3
                    let mut out = if let Some(q4k_bytes0) = q4k_layer
    1276
    3
                        .and_then(|q| q.ffn_down_weight0.as_ref0())
    1277
                    {
    1278
    0
                        let mut output = Vec::with_capacity(seq_len * hidden_dim);
    1279
    0
                        for s in 0..seq_len {
    1280
    0
                            let input_slice = &ffn_hidden[s * intermediate_dim..(s + 1) * intermediate_dim];
    1281
    0
                            let pos_out = matmul_q4k_rowmajor(q4k_bytes, input_slice, hidden_dim, intermediate_dim);
    1282
    0
                            output.extend(pos_out);
    1283
    0
                        }
    1284
    0
                        output
    1285
                    } else {
    1286
    3
                        self.matmul(&ffn_hidden, &layer.ffn_down_weight, intermediate_dim, hidden_dim)
    1287
                    };
    1288
    3
                    if let Some(ref bias0) = layer.ffn_down_bias {
    1289
    0
                        self.add_bias(&mut out, bias);
    1290
    3
                    }
    1291
    3
                    out
    1292
                };
    1293
    1294
                // 2h. Residual connection
    1295
    36
                for i in 0..hidden3.len3() {
    1296
    36
                    hidden[i] += ffn_output[i];
    1297
    36
                }
    1298
            }
    1299
    1300
            // 3. Final layer norm
    1301
    3
            let normed = self.layer_norm(
    1302
    3
                &hidden,
    1303
    3
                &self.output_norm_weight,
    1304
    3
                self.output_norm_bias.as_deref(),
    1305
    3
                self.config.eps,
    1306
            );
    1307
    1308
            // 4. LM head projection (only last token)
    1309
    3
            let seq_len = token_ids.len();
    1310
    3
            let last_hidden_start = (seq_len - 1) * hidden_dim;
    1311
    3
            let last_hidden = &normed[last_hidden_start..last_hidden_start + hidden_dim];
    1312
    1313
    3
            let mut logits = self.matmul(
    1314
    3
                last_hidden,
    1315
    3
                &self.lm_head_weight,
    1316
    3
                hidden_dim,
    1317
    3
                self.config.vocab_size,
    1318
            );
    1319
    3
            if let Some(ref bias0) = self.lm_head_bias {
    1320
    0
                self.add_bias(&mut logits, bias);
    1321
    3
            }
    1322
    1323
    3
            Ok(logits)
    1324
    3
        }
    1325
    1326
        /// Predict next token (greedy decoding)
    1327
        ///
    1328
        /// # Arguments
    1329
        ///
    1330
        /// * `token_ids` - Input token IDs
    1331
        ///
    1332
        /// # Returns
    1333
        ///
    1334
        /// Token ID with highest probability
    1335
        ///
    1336
        /// # Errors
    1337
        ///
    1338
        /// Returns error if inference fails
    1339
    0
        pub fn predict_next(&self, token_ids: &[u32]) -> Result<u32> {
    1340
    0
            let logits = self.forward(token_ids)?;
    1341
    1342
            // Argmax
    1343
    0
            let (max_idx, _) = logits
    1344
    0
                .iter()
    1345
    0
                .enumerate()
    1346
    0
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    1347
    0
                .ok_or_else(|| RealizarError::InvalidShape {
    1348
    0
                    reason: "Empty logits".to_string(),
    1349
    0
                })?;
    1350
    1351
    0
            Ok(max_idx as u32)
    1352
    0
        }
    1353
    1354
        /// Forward pass with KV cache for efficient autoregressive generation (Y4)
    1355
        ///
    1356
        /// Processes a single token using cached key-value pairs from previous positions.
    1357
        ///
    1358
        /// # Arguments
    1359
        ///
    1360
        /// * `token_id` - Single token ID to process
    1361
        /// * `cache` - Mutable KV cache to read from and append to
    1362
        /// * `position` - Position in sequence (0-indexed)
    1363
        ///
    1364
        /// # Returns
    1365
        ///
    1366
        /// Logits over vocabulary for next token prediction
    1367
    0
        pub fn forward_with_cache(
    1368
    0
            &self,
    1369
    0
            token_id: u32,
    1370
    0
            cache: &mut AprKVCache,
    1371
    0
            position: usize,
    1372
    0
        ) -> Result<Vec<f32>> {
    1373
            // DEBUG: Force F32 fallback to verify data layout issues
    1374
    0
            let force_f32 = std::env::var("APR_FORCE_F32").is_ok();
    1375
    1376
    0
            let hidden_dim = self.config.hidden_dim;
    1377
    0
            let num_heads = self.config.num_heads;
    1378
    0
            let num_kv_heads = self.config.num_kv_heads;
    1379
    0
            let head_dim = hidden_dim / num_heads;
    1380
    1381
            // 1. Token embedding lookup
    1382
    0
            let mut hidden = self.embed(&[token_id]);
    1383
    1384
            // 2. Process through transformer layers
    1385
    0
            let layers_start = std::time::Instant::now();
    1386
    0
            for (layer_idx, layer) in self.layers.iter().enumerate() {
    1387
                // PMAT-103: Get Q4K weights for this layer (if available) for fused kernels
    1388
    0
                let q4k_layer = self.q4k_layers.as_ref().and_then(|l| l.get(layer_idx));
    1389
    1390
                // 2a. Attention layer norm
    1391
    0
                let normed = self.layer_norm(
    1392
    0
                    &hidden,
    1393
    0
                    &layer.attn_norm_weight,
    1394
    0
                    layer.attn_norm_bias.as_deref(),
    1395
    0
                    self.config.eps,
    1396
                );
    1397
    1398
                // 2b. QKV projection (single token)
    1399
                // PMAT-103: Use fused Q4K kernels for separate Q, K, V weights when available
    1400
    0
                let kv_size = num_kv_heads * head_dim;
    1401
    0
                let (mut q, mut k, v) = if let Some(q4k) = q4k_layer {
    1402
                    // Try Q4K fused kernels for Q, K
    1403
    0
                    let q = if let Some(ref q_bytes) = q4k.attn_q_weight {
    1404
    0
                        if layer_idx == 0 && position == 0 {
    1405
    0
                            eprintln!("[TRACE-CACHE] Layer 0: Q projection using Q4K fused kernel");
    1406
    0
                        }
    1407
    0
                        matmul_q4k_rowmajor(q_bytes, &normed, hidden_dim, hidden_dim)
    1408
                    } else {
    1409
                        // Fallback to F32 for Q (should not happen for GGUF models)
    1410
    0
                        let q_weight = &layer.qkv_weight[0..hidden_dim * hidden_dim];
    1411
    0
                        self.matmul(&normed, q_weight, hidden_dim, hidden_dim)
    1412
                    };
    1413
    1414
    0
                    let k = if let Some(ref k_bytes) = q4k.attn_k_weight {
    1415
    0
                        if layer_idx == 0 && position == 0 {
    1416
    0
                            eprintln!("[TRACE-CACHE] Layer 0: K projection using Q4K fused kernel");
    1417
    0
                        }
    1418
    0
                        matmul_q4k_rowmajor(k_bytes, &normed, kv_size, hidden_dim)
    1419
                    } else {
    1420
    0
                        let k_start = hidden_dim * hidden_dim;
    1421
    0
                        let k_weight = &layer.qkv_weight[k_start..k_start + kv_size * hidden_dim];
    1422
    0
                        self.matmul(&normed, k_weight, hidden_dim, kv_size)
    1423
                    };
    1424
    1425
                    // V can be Q4K or Q6K
    1426
    0
                    let v = if let Some(ref v_bytes) = q4k.attn_v_weight {
    1427
    0
                        if layer_idx == 0 && position == 0 {
    1428
    0
                            eprintln!("[TRACE-CACHE] Layer 0: V projection using Q4K fused kernel");
    1429
    0
                        }
    1430
    0
                        matmul_q4k_rowmajor(v_bytes, &normed, kv_size, hidden_dim)
    1431
    0
                    } else if let Some(ref v_bytes) = q4k.attn_v_weight_q6k {
    1432
    0
                        if layer_idx == 0 && position == 0 {
    1433
    0
                            eprintln!("[TRACE-CACHE] Layer 0: V projection using Q6K fused kernel");
    1434
    0
                        }
    1435
    0
                        matmul_q6k_rowmajor(v_bytes, &normed, kv_size, hidden_dim)
    1436
                    } else {
    1437
    0
                        let v_start = hidden_dim * hidden_dim + kv_size * hidden_dim;
    1438
    0
                        let v_weight = &layer.qkv_weight[v_start..v_start + kv_size * hidden_dim];
    1439
    0
                        self.matmul(&normed, v_weight, hidden_dim, kv_size)
    1440
                    };
    1441
    1442
    0
                    (q, k, v)
    1443
                } else {
    1444
                    // Fallback: Combined QKV with F32 (legacy path)
    1445
    0
                    if layer_idx == 0 && position == 0 {
    1446
    0
                        eprintln!("[TRACE-CACHE] Layer 0: QKV projection using F32 (not fused)");
    1447
    0
                    }
    1448
    0
                    let qkv_out_dim = layer.qkv_weight.len() / hidden_dim;
    1449
    0
                    let mut qkv = self.matmul(&normed, &layer.qkv_weight, hidden_dim, qkv_out_dim);
    1450
    0
                    if let Some(ref bias) = layer.qkv_bias {
    1451
    0
                        self.add_bias(&mut qkv, bias);
    1452
    0
                    }
    1453
    0
                    let q = qkv[0..hidden_dim].to_vec();
    1454
    0
                    let k = qkv[hidden_dim..hidden_dim + kv_size].to_vec();
    1455
    0
                    let v = qkv[hidden_dim + kv_size..hidden_dim + 2 * kv_size].to_vec();
    1456
    0
                    (q, k, v)
    1457
                };
    1458
    1459
                // Apply biases if present (for fused path)
    1460
                // The combined qkv_bias is [Q_bias | K_bias | V_bias]
    1461
    0
                let mut v_mut = v;
    1462
    0
                if q4k_layer.is_some() {
    1463
    0
                    if let Some(ref bias) = layer.qkv_bias {
    1464
                        // Split bias into Q, K, V portions
    1465
    0
                        for (i, b) in bias[0..hidden_dim].iter().enumerate() {
    1466
    0
                            q[i] += b;
    1467
    0
                        }
    1468
    0
                        for (i, b) in bias[hidden_dim..hidden_dim + kv_size].iter().enumerate() {
    1469
    0
                            k[i] += b;
    1470
    0
                        }
    1471
                        // V bias starts after Q and K biases
    1472
    0
                        let v_bias_start = hidden_dim + kv_size;
    1473
    0
                        for (i, b) in bias[v_bias_start..v_bias_start + kv_size].iter().enumerate() {
    1474
    0
                            v_mut[i] += b;
    1475
    0
                        }
    1476
    0
                    }
    1477
    0
                }
    1478
    0
                let v = v_mut;
    1479
    1480
                // PMAT-103: Apply RoPE to Q and K at current position
    1481
                // This was missing, causing garbage output
    1482
    0
                self.apply_rope_f32(&mut q, position, num_heads, head_dim);
    1483
    0
                self.apply_rope_f32(&mut k, position, num_kv_heads, head_dim);
    1484
    1485
                // 2c. Append K, V to cache (K now has RoPE applied)
    1486
    0
                cache.append(layer_idx, &k, &v);
    1487
    1488
                // 2d. Compute attention with full cache
    1489
    0
                let (k_cache, v_cache) = cache.get(layer_idx);
    1490
    0
                let seq_len = cache.len();
    1491
    1492
                // Simplified attention: compute Q·K^T / sqrt(d), softmax, then V
    1493
    0
                let mut attn_out = vec![0.0f32; hidden_dim];
    1494
    1495
                // PMAT-103: SIMD-accelerated attention computation
    1496
    0
                let scale = 1.0 / (head_dim as f32).sqrt();
    1497
    1498
    0
                for h in 0..num_heads {
    1499
    0
                    let kv_head = h * num_kv_heads / num_heads; // GQA mapping
    1500
    0
                    let q_start = h * head_dim;
    1501
    0
                    let q_slice = &q[q_start..q_start + head_dim]; // q is now Vec with RoPE applied
    1502
    1503
                    // Compute attention scores with SIMD dot product
    1504
    0
                    let mut scores = Vec::with_capacity(seq_len);
    1505
    0
                    for pos in 0..seq_len {
    1506
    0
                        let k_start = pos * kv_size + kv_head * head_dim;
    1507
    0
                        let k_slice = &k_cache[k_start..k_start + head_dim];
    1508
    0
                        // SIMD dot product (AVX2 when available)
    1509
    0
                        let dot = simd_dot_f32(q_slice, k_slice);
    1510
    0
                        scores.push(dot * scale);
    1511
    0
                    }
    1512
    1513
                    // Causal mask: only attend to positions <= current
    1514
    0
                    for pos in (position + 1)..seq_len {
    1515
    0
                        scores[pos] = f32::NEG_INFINITY;
    1516
    0
                    }
    1517
    1518
                    // Softmax (scalar - typically small seq_len during decode)
    1519
    0
                    let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1520
    0
                    let mut exp_scores: Vec<f32> =
    1521
    0
                        scores.iter().map(|s| (s - max_score).exp()).collect();
    1522
    0
                    let sum: f32 = exp_scores.iter().sum();
    1523
    0
                    if sum > 0.0 {
    1524
    0
                        let inv_sum = 1.0 / sum;
    1525
    0
                        for s in &mut exp_scores {
    1526
    0
                            *s *= inv_sum;
    1527
    0
                        }
    1528
    0
                    }
    1529
    1530
                    // Weighted sum of V with SIMD accumulation
    1531
    0
                    let attn_out_head = &mut attn_out[q_start..q_start + head_dim];
    1532
    0
                    for pos in 0..seq_len {
    1533
    0
                        let v_start = pos * kv_size + kv_head * head_dim;
    1534
    0
                        let v_slice = &v_cache[v_start..v_start + head_dim];
    1535
    0
                        // SIMD weighted accumulation (AVX2 when available)
    1536
    0
                        simd_add_weighted(attn_out_head, v_slice, exp_scores[pos]);
    1537
    0
                    }
    1538
                }
    1539
    1540
                // 2e. Attention output projection
    1541
                // PMAT-103: Use Q4K fused kernel when available (single token path)
    1542
    0
                let mut attn_output = if !force_f32 {
    1543
    0
                    if let Some(q4k_bytes) = q4k_layer
    1544
    0
                        .and_then(|q| q.attn_output_weight.as_ref())
    1545
                    {
    1546
    0
                        if layer_idx == 0 && position == 0 {
    1547
    0
                            eprintln!("[TRACE-CACHE] Layer 0: attn_output using Q4K fused kernel");
    1548
    0
                        }
    1549
    0
                        matmul_q4k_rowmajor(q4k_bytes, &attn_out, hidden_dim, hidden_dim)
    1550
                    } else {
    1551
    0
                        if layer_idx == 0 && position == 0 {
    1552
    0
                            eprintln!("[TRACE-CACHE] Layer 0: attn_output using F32 fallback (slow!)");
    1553
    0
                        }
    1554
    0
                        self.matmul(&attn_out, &layer.attn_output_weight, hidden_dim, hidden_dim)
    1555
                    }
    1556
                } else {
    1557
    0
                    if layer_idx == 0 && position == 0 {
    1558
    0
                        eprintln!("[TRACE-CACHE] Layer 0: attn_output using F32 (APR_FORCE_F32)");
    1559
    0
                    }
    1560
    0
                    self.matmul(&attn_out, &layer.attn_output_weight, hidden_dim, hidden_dim)
    1561
                };
    1562
    0
                if let Some(ref bias) = layer.attn_output_bias {
    1563
    0
                    self.add_bias(&mut attn_output, bias);
    1564
    0
                }
    1565
    1566
                // 2f. Residual connection
    1567
    0
                for i in 0..hidden.len() {
    1568
    0
                    hidden[i] += attn_output[i];
    1569
    0
                }
    1570
    1571
                // 2g. Apply FFN norm if present (post_attention_layernorm)
    1572
    0
                let ffn_input = if let Some(ref ffn_norm) = layer.ffn_norm_weight {
    1573
    0
                    self.layer_norm(
    1574
    0
                        &hidden,
    1575
    0
                        ffn_norm,
    1576
    0
                        layer.ffn_norm_bias.as_deref(),
    1577
    0
                        self.config.eps,
    1578
                    )
    1579
                } else {
    1580
    0
                    hidden.clone()
    1581
                };
    1582
    1583
                // 2h. FFN projection (SwiGLU or standard GELU)
    1584
                // PMAT-103 FIX: Use Q4K/Q6K fused kernels when available (single token path)
    1585
    0
                let intermediate_dim = self.config.intermediate_dim;
    1586
    0
                let ffn_output = if let Some(ref _gate_weight) = layer.ffn_gate_weight {
    1587
                    // SwiGLU: down(SiLU(gate(x)) * up(x))
    1588
                    // PMAT-103: Check for Q4K gate weight
    1589
    0
                    let gate = if !force_f32 {
    1590
    0
                        if let Some(q4k_bytes) = q4k_layer
    1591
    0
                            .and_then(|q| q.ffn_gate_weight.as_ref())
    1592
                        {
    1593
    0
                            if layer_idx == 0 && position == 0 {
    1594
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_gate using Q4K fused kernel");
    1595
    0
                            }
    1596
    0
                            matmul_q4k_rowmajor(q4k_bytes, &ffn_input, intermediate_dim, hidden_dim)
    1597
                        } else {
    1598
    0
                            if layer_idx == 0 && position == 0 {
    1599
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_gate using F32 fallback (slow!)");
    1600
    0
                            }
    1601
    0
                            self.matmul(&ffn_input, layer.ffn_gate_weight.as_ref().expect("gate weight"), hidden_dim, intermediate_dim)
    1602
                        }
    1603
                    } else {
    1604
    0
                        if layer_idx == 0 && position == 0 {
    1605
    0
                            eprintln!("[TRACE-CACHE] Layer 0: ffn_gate using F32 (APR_FORCE_F32)");
    1606
    0
                        }
    1607
    0
                        self.matmul(&ffn_input, layer.ffn_gate_weight.as_ref().expect("gate weight"), hidden_dim, intermediate_dim)
    1608
                    };
    1609
    1610
                    // PMAT-103: Check for Q4K/Q6K up weight
    1611
    0
                    let up = if !force_f32 {
    1612
    0
                        if let Some(q4k_bytes) = q4k_layer
    1613
    0
                            .and_then(|q| q.ffn_up_weight.as_ref())
    1614
                        {
    1615
    0
                            if layer_idx == 0 && position == 0 {
    1616
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_up using Q4K fused kernel");
    1617
    0
                            }
    1618
    0
                            matmul_q4k_rowmajor(q4k_bytes, &ffn_input, intermediate_dim, hidden_dim)
    1619
    0
                        } else if let Some(q6k_bytes) = q4k_layer
    1620
    0
                            .and_then(|q| q.ffn_up_weight_q6k.as_ref())
    1621
                        {
    1622
    0
                            if layer_idx == 0 && position == 0 {
    1623
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_up using Q6K fused kernel");
    1624
    0
                            }
    1625
    0
                            matmul_q6k_rowmajor(q6k_bytes, &ffn_input, intermediate_dim, hidden_dim)
    1626
                        } else {
    1627
    0
                            if layer_idx == 0 && position == 0 {
    1628
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_up using F32 fallback (slow!)");
    1629
    0
                            }
    1630
    0
                            self.matmul(&ffn_input, &layer.ffn_up_weight, hidden_dim, intermediate_dim)
    1631
                        }
    1632
                    } else {
    1633
    0
                        if layer_idx == 0 && position == 0 {
    1634
    0
                            eprintln!("[TRACE-CACHE] Layer 0: ffn_up using F32 (APR_FORCE_F32)");
    1635
    0
                        }
    1636
    0
                        self.matmul(&ffn_input, &layer.ffn_up_weight, hidden_dim, intermediate_dim)
    1637
                    };
    1638
    1639
                    // SiLU(gate) * up, then down projection
    1640
    0
                    let mut ffn_hidden = Vec::with_capacity(gate.len());
    1641
    0
                    for (g, u) in gate.iter().zip(up.iter()) {
    1642
    0
                        let silu_g = g / (1.0 + (-g).exp()); // SiLU = x * sigmoid(x)
    1643
    0
                        ffn_hidden.push(silu_g * u);
    1644
    0
                    }
    1645
    1646
                    // PMAT-103: Check for Q4K or Q6K down weight
    1647
    0
                    let mut out = if !force_f32 {
    1648
    0
                        if let Some(q4k_bytes) = q4k_layer
    1649
    0
                            .and_then(|q| q.ffn_down_weight.as_ref())
    1650
                        {
    1651
    0
                            if layer_idx == 0 && position == 0 {
    1652
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_down using Q4K fused kernel");
    1653
    0
                            }
    1654
    0
                            matmul_q4k_rowmajor(q4k_bytes, &ffn_hidden, hidden_dim, intermediate_dim)
    1655
    0
                        } else if let Some(q6k_bytes) = q4k_layer
    1656
    0
                            .and_then(|q| q.ffn_down_weight_q6k.as_ref())
    1657
                        {
    1658
    0
                            if layer_idx == 0 && position == 0 {
    1659
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_down using Q6K fused kernel");
    1660
    0
                            }
    1661
    0
                            matmul_q6k_rowmajor(q6k_bytes, &ffn_hidden, hidden_dim, intermediate_dim)
    1662
                        } else {
    1663
    0
                            if layer_idx == 0 && position == 0 {
    1664
    0
                                eprintln!("[TRACE-CACHE] Layer 0: ffn_down using F32 fallback (slow!)");
    1665
    0
                            }
    1666
    0
                            self.matmul(&ffn_hidden, &layer.ffn_down_weight, intermediate_dim, hidden_dim)
    1667
                        }
    1668
                    } else {
    1669
    0
                        if layer_idx == 0 && position == 0 {
    1670
    0
                            eprintln!("[TRACE-CACHE] Layer 0: ffn_down using F32 (APR_FORCE_F32)");
    1671
    0
                        }
    1672
    0
                        self.matmul(&ffn_hidden, &layer.ffn_down_weight, intermediate_dim, hidden_dim)
    1673
                    };
    1674
    0
                    if let Some(ref bias) = layer.ffn_down_bias {
    1675
    0
                        self.add_bias(&mut out, bias);
    1676
    0
                    }
    1677
    0
                    out
    1678
                } else {
    1679
                    // Standard MLP: down(GELU(up(x)))
    1680
                    // PMAT-103: Check for Q4K up weight
    1681
    0
                    let mut ffn_hidden = if let Some(q4k_bytes) = q4k_layer
    1682
    0
                        .and_then(|q| q.ffn_up_weight.as_ref())
    1683
                    {
    1684
    0
                        matmul_q4k_rowmajor(q4k_bytes, &ffn_input, intermediate_dim, hidden_dim)
    1685
                    } else {
    1686
    0
                        self.matmul(&ffn_input, &layer.ffn_up_weight, hidden_dim, intermediate_dim)
    1687
                    };
    1688
    0
                    if let Some(ref bias) = layer.ffn_up_bias {
    1689
    0
                        self.add_bias(&mut ffn_hidden, bias);
    1690
    0
                    }
    1691
    0
                    self.gelu(&mut ffn_hidden);
    1692
    1693
                    // PMAT-103: Check for Q4K down weight
    1694
    0
                    let mut out = if let Some(q4k_bytes) = q4k_layer
    1695
    0
                        .and_then(|q| q.ffn_down_weight.as_ref())
    1696
                    {
    1697
    0
                        matmul_q4k_rowmajor(q4k_bytes, &ffn_hidden, hidden_dim, intermediate_dim)
    1698
                    } else {
    1699
    0
                        self.matmul(&ffn_hidden, &layer.ffn_down_weight, intermediate_dim, hidden_dim)
    1700
                    };
    1701
    0
                    if let Some(ref bias) = layer.ffn_down_bias {
    1702
    0
                        self.add_bias(&mut out, bias);
    1703
    0
                    }
    1704
    0
                    out
    1705
                };
    1706
    1707
                // 2i. Residual connection
    1708
    0
                for i in 0..hidden.len() {
    1709
    0
                    hidden[i] += ffn_output[i];
    1710
    0
                }
    1711
            }
    1712
    0
            eprintln!("[TRACE-CACHE] pos={}: {} layers took {:?}", position, self.layers.len(), layers_start.elapsed());
    1713
    1714
            // 3. Final layer norm
    1715
    0
            let normed = self.layer_norm(
    1716
    0
                &hidden,
    1717
    0
                &self.output_norm_weight,
    1718
    0
                self.output_norm_bias.as_deref(),
    1719
    0
                self.config.eps,
    1720
            );
    1721
    1722
            // 4. LM head projection
    1723
            // PMAT-103: Use Q4K/Q6K fused kernel when available (single token path)
    1724
    0
            let lm_start = std::time::Instant::now();
    1725
    0
            let mut logits = if !force_f32 {
    1726
    0
                if let Some(ref q4k_bytes) = self.lm_head_weight_q4k {
    1727
    0
                    eprintln!("[TRACE-CACHE] lm_head using Q4K fused kernel");
    1728
    0
                    matmul_q4k_rowmajor(q4k_bytes, &normed, self.config.vocab_size, hidden_dim)
    1729
    0
                } else if let Some(ref q6k_bytes) = self.lm_head_weight_q6k {
    1730
    0
                    let result = matmul_q6k_rowmajor(q6k_bytes, &normed, self.config.vocab_size, hidden_dim);
    1731
    0
                    eprintln!("[TRACE-CACHE] lm_head Q6K took {:?}", lm_start.elapsed());
    1732
    0
                    result
    1733
                } else {
    1734
    0
                    self.matmul(
    1735
    0
                        &normed,
    1736
    0
                        &self.lm_head_weight,
    1737
    0
                        hidden_dim,
    1738
    0
                        self.config.vocab_size,
    1739
                    )
    1740
                }
    1741
            } else {
    1742
    0
                eprintln!("[TRACE-CACHE] lm_head using F32 (APR_FORCE_F32)");
    1743
    0
                self.matmul(
    1744
    0
                    &normed,
    1745
    0
                    &self.lm_head_weight,
    1746
    0
                    hidden_dim,
    1747
    0
                    self.config.vocab_size,
    1748
                )
    1749
            };
    1750
    0
            if let Some(ref bias) = self.lm_head_bias {
    1751
    0
                self.add_bias(&mut logits, bias);
    1752
    0
            }
    1753
    1754
    0
            Ok(logits)
    1755
    0
        }
    1756
    1757
        /// Generate tokens using KV cache for efficiency (Y4)
    1758
        ///
    1759
        /// # Arguments
    1760
        ///
    1761
        /// * `prompt` - Initial token IDs
    1762
        /// * `config` - Generation configuration
    1763
        ///
    1764
        /// # Returns
    1765
        ///
    1766
        /// Generated token sequence (including prompt)
    1767
    0
        pub fn generate_with_cache(&self, prompt: &[u32], config: &GenerateConfig) -> Result<Vec<u32>> {
    1768
    0
            if prompt.is_empty() {
    1769
    0
                return Err(RealizarError::InvalidShape {
    1770
    0
                    reason: "Prompt cannot be empty".to_string(),
    1771
    0
                });
    1772
    0
            }
    1773
    1774
    0
            let mut cache = AprKVCache::new(&self.config);
    1775
    0
            let mut output = prompt.to_vec();
    1776
    1777
            // PMAT-103 FIX: Process prompt tokens and KEEP the logits from the last one.
    1778
            // Previously we threw away all logits (`let _ = ...`) and then reprocessed
    1779
            // the last prompt token at the same position, corrupting the KV cache.
    1780
    0
            let mut logits = Vec::new();
    1781
    1782
            // PMAT-103 TRACE: Measure per-token timing to verify O(n) vs O(n²)
    1783
    0
            let trace_enabled = std::env::var("REALIZE_TRACE").is_ok();
    1784
    0
            if trace_enabled {
    1785
    0
                eprintln!("[TRACE] Processing {} prompt tokens...", prompt.len());
    1786
    0
            }
    1787
    1788
    0
            for (pos, &token) in prompt.iter().enumerate() {
    1789
    0
                let start = std::time::Instant::now();
    1790
    0
                logits = self.forward_with_cache(token, &mut cache, pos)?;
    1791
    0
                if trace_enabled {
    1792
    0
                    eprintln!("[TRACE] Prompt token {}: {:?}", pos, start.elapsed());
    1793
    0
                }
    1794
            }
    1795
    1796
            // Generate new tokens using the logits we already have
    1797
    0
            for i in 0..config.max_tokens {
    1798
                // Sample from current logits (which predict the NEXT token)
    1799
    0
                let next_token = if config.temperature == 0.0 {
    1800
    0
                    logits
    1801
    0
                        .iter()
    1802
    0
                        .enumerate()
    1803
    0
                        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    1804
    0
                        .map_or(0, |(idx, _)| idx as u32)
    1805
                } else {
    1806
    0
                    let scaled: Vec<f32> = logits.iter().map(|l| l / config.temperature).collect();
    1807
    0
                    let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    1808
    0
                    let exp_vals: Vec<f32> = scaled.iter().map(|s| (s - max_val).exp()).collect();
    1809
    0
                    let sum: f32 = exp_vals.iter().sum();
    1810
    0
                    let probs: Vec<f32> = exp_vals.iter().map(|e| e / sum).collect();
    1811
    0
                    probs
    1812
    0
                        .iter()
    1813
    0
                        .enumerate()
    1814
    0
                        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    1815
    0
                        .map_or(0, |(idx, _)| idx as u32)
    1816
                };
    1817
    1818
    0
                output.push(next_token);
    1819
    1820
                // Check for EOS tokens
    1821
    0
                if next_token == 0 || next_token == 2 || next_token == 151645 || next_token == 151643 {
    1822
    0
                    break;
    1823
    0
                }
    1824
    1825
                // If we need more tokens, process this one to get logits for the next
    1826
    0
                if i < config.max_tokens - 1 {
    1827
                    // Position is output.len() - 1 = prompt.len() + (i + 1) - 1 = prompt.len() + i
    1828
    0
                    let start = std::time::Instant::now();
    1829
    0
                    logits = self.forward_with_cache(next_token, &mut cache, output.len() - 1)?;
    1830
    0
                    if trace_enabled {
    1831
    0
                        eprintln!("[TRACE] Gen token {} (pos {}): {:?}", i, output.len() - 1, start.elapsed());
    1832
    0
                    }
    1833
    0
                }
    1834
            }
    1835
    1836
    0
            if trace_enabled {
    1837
    0
                eprintln!("[TRACE] Generation complete. Total output tokens: {}", output.len());
    1838
    0
            }
    1839
    1840
    0
            Ok(output)
    1841
    0
        }
    1842
    }
    1843
    1844
    /// Convert from `GGUFTransformer` to APR format
    1845
    ///
    1846
    /// This dequantizes all GGUF weights to F32 for WASM compatibility.
    1847
    #[cfg(feature = "default")]
    1848
    impl From<&crate::gguf::GGUFTransformer> for AprTransformer {
    1849
    0
        fn from(gguf: &crate::gguf::GGUFTransformer) -> Self {
    1850
    0
            let config = AprTransformerConfig {
    1851
    0
                architecture: gguf.config.architecture.clone(),
    1852
    0
                hidden_dim: gguf.config.hidden_dim,
    1853
    0
                num_layers: gguf.config.num_layers,
    1854
    0
                num_heads: gguf.config.num_heads,
    1855
    0
                num_kv_heads: gguf.config.num_kv_heads,
    1856
    0
                vocab_size: gguf.config.vocab_size,
    1857
    0
                intermediate_dim: gguf.config.intermediate_dim,
    1858
    0
                context_length: gguf.config.context_length,
    1859
    0
                rope_theta: gguf.config.rope_theta,
    1860
    0
                eps: gguf.config.eps,
    1861
    0
            };
    1862
    1863
    0
            let layers = gguf
    1864
    0
                .layers
    1865
    0
                .iter()
    1866
    0
                .map(|l| AprTransformerLayer {
    1867
    0
                    attn_norm_weight: l.attn_norm_weight.clone(),
    1868
    0
                    attn_norm_bias: l.attn_norm_bias.clone(),
    1869
    0
                    qkv_weight: l.qkv_weight.clone(),
    1870
    0
                    qkv_bias: l.qkv_bias.clone(),
    1871
    0
                    attn_output_weight: l.attn_output_weight.clone(),
    1872
    0
                    attn_output_bias: l.attn_output_bias.clone(),
    1873
    0
                    ffn_gate_weight: l.ffn_gate_weight.clone(),
    1874
    0
                    ffn_gate_bias: l.ffn_gate_bias.clone(),
    1875
    0
                    ffn_up_weight: l.ffn_up_weight.clone(),
    1876
    0
                    ffn_up_bias: l.ffn_up_bias.clone(),
    1877
    0
                    ffn_down_weight: l.ffn_down_weight.clone(),
    1878
    0
                    ffn_down_bias: l.ffn_down_bias.clone(),
    1879
    0
                    ffn_norm_weight: l.ffn_norm_weight.clone(),
    1880
    0
                    ffn_norm_bias: l.ffn_norm_bias.clone(),
    1881
    0
                })
    1882
    0
                .collect();
    1883
    1884
    0
            Self {
    1885
    0
                config,
    1886
    0
                token_embedding: gguf.token_embedding.clone(),
    1887
    0
                layers,
    1888
    0
                output_norm_weight: gguf.output_norm_weight.clone(),
    1889
    0
                output_norm_bias: gguf.output_norm_bias.clone(),
    1890
    0
                lm_head_weight: gguf.lm_head_weight.clone(),
    1891
    0
                lm_head_bias: gguf.lm_head_bias.clone(),
    1892
    0
                q4k_layers: None,
    1893
    0
                lm_head_weight_q6k: None,
    1894
    0
                lm_head_weight_q4k: None,
    1895
    0
            }
    1896
    0
        }
    1897
    }
    1898
    // Tests shattered to tests/ directory (PMAT-803)
    1899
    #[cfg(test)]
    1900
    mod tests;