## Coverage Report

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

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

    /home/noah/src/trueno/src/backends/q4k/mod.rs

    Line
    Count
    Source
    1
    //! Fused Q4_K Matrix-Vector Multiply (F-GPU-130)
    2
    //!
    3
    //! This module implements fused quantized matrix-vector multiplication that operates
    4
    //! directly on Q4_K compressed weights without full dequantization.
    5
    //!
    6
    //! # Q4_K Format (llama.cpp compatible)
    7
    //!
    8
    //! Super-block layout (144 bytes per 256 elements):
    9
    //! - `d`: 2 bytes (f16 global scale)
    10
    //! - `dmin`: 2 bytes (f16 global min scale)
    11
    //! - `scales`: 12 bytes (packed 6-bit scales and mins for 8 sub-blocks)
    12
    //! - `qs`: 128 bytes (4-bit quantized values, interleaved low/high nibbles)
    13
    //!
    14
    //! # Golden Test Invariant (Section 12.4 of spec)
    15
    //!
    16
    //! For all Q4K weight W and input x:
    17
    //! ```text
    18
    //! matmul_q4k_f32(W, x) ≈ matmul(dequant_q4k_to_f32(W), x)  within ε = 1e-3
    19
    //! ```
    20
    //!
    21
    //! # Performance Targets
    22
    //!
    23
    //! - Baseline (dequant+matmul): 0.27 tok/s
    24
    //! - Target (fused): >5 tok/s CPU, >100 tok/s GPU
    25
    //!
    26
    //! # Example
    27
    //!
    28
    //! ```rust,ignore
    29
    //! use trueno::backends::q4k::matmul_q4k_f32;
    30
    //!
    31
    //! let q4k_weights = load_q4k_tensor("gate_proj.weight");
    32
    //! let input = vec![1.0f32; 896];
    33
    //! let output = matmul_q4k_f32(&q4k_weights, &input, 4864, 896);
    34
    //! ```
    35
    36
    #![allow(dead_code)]
    37
    38
    // Sub-modules
    39
    mod colmajor;
    40
    mod dequant;
    41
    mod gemv;
    42
    43
    // Re-exports
    44
    pub use colmajor::{matmul_q4k_f32_colmajor, matmul_q4k_f32_colmajor_dispatch};
    45
    pub use dequant::dequantize_q4k_to_f32;
    46
    pub use gemv::{matmul_q4k_f32, matmul_q4k_f32_dispatch, matmul_q4k_f32_scalar};
    47
    48
    // Constants (pub(crate) for submodule access)
    49
    pub(crate) const SUPER_BLOCK_SIZE: usize = 256;
    50
    pub(crate) const SUPER_BLOCK_BYTES: usize = 144;
    51
    #[allow(dead_code)] // Reserved for future sub-block optimizations
    52
    pub(crate) const SUB_BLOCK_SIZE: usize = 32;
    53
    54
    /// Convert f16 bits to f32
    55
    #[inline(always)]
    56
    0
    fn f16_to_f32(bits: u16) -> f32 {
    57
    0
        let sign = ((bits & 0x8000) as u32) << 16;
    58
    0
        let exp = (bits >> 10) & 0x1F;
    59
    0
        let mantissa = (bits & 0x3FF) as u32;
    60
    61
    0
        if exp == 0 {
    62
    0
            if mantissa == 0 {
    63
    0
                f32::from_bits(sign)
    64
            } else {
    65
                // Subnormal
    66
    0
                let mut m = mantissa;
    67
    0
                let mut e = 0i32;
    68
    0
                while (m & 0x400) == 0 {
    69
    0
                    m <<= 1;
    70
    0
                    e -= 1;
    71
    0
                }
    72
    0
                let new_exp = ((127 - 15 + 1 + e) as u32) << 23;
    73
    0
                let new_mantissa = (m & 0x3FF) << 13;
    74
    0
                f32::from_bits(sign | new_exp | new_mantissa)
    75
            }
    76
    0
        } else if exp == 31 {
    77
    0
            f32::from_bits(sign | (0xFF << 23) | (mantissa << 13))
    78
        } else {
    79
    0
            let new_exp = ((exp as i32 - 15 + 127) as u32) << 23;
    80
    0
            f32::from_bits(sign | new_exp | (mantissa << 13))
    81
        }
    82
    0
    }
    83
    84
    /// Parse Q4_K super-block header and scales
    85
    ///
    86
    /// Returns (d, dmin, scales[8], mins[8])
    87
    #[inline(always)]
    88
    0
    pub(crate) fn parse_q4k_header(block: &[u8]) -> (f32, f32, [u8; 8], [u8; 8]) {
    89
    0
        debug_assert!(block.len() >= 16);
    90
    91
        // Read d and dmin (f16)
    92
    0
        let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]]));
    93
    0
        let dmin = f16_to_f32(u16::from_le_bytes([block[2], block[3]]));
    94
    95
        // Unpack scales and mins (llama.cpp format)
    96
    0
        let scales_bytes = &block[4..16];
    97
    0
        let mut scales = [0u8; 8];
    98
    0
        let mut mins = [0u8; 8];
    99
    100
    0
        for i in 0..4 {
    101
    0
            // Blocks 0-3: lower 6 bits of bytes 0-3 and 4-7
    102
    0
            scales[i] = scales_bytes[i] & 0x3F;
    103
    0
            mins[i] = scales_bytes[i + 4] & 0x3F;
    104
    0
            // Blocks 4-7: lower 4 bits from bytes 8-11, upper 2 bits from bytes 0-3/4-7
    105
    0
            scales[i + 4] = (scales_bytes[i + 8] & 0x0F) | ((scales_bytes[i] >> 6) << 4);
    106
    0
            mins[i + 4] = (scales_bytes[i + 8] >> 4) | ((scales_bytes[i + 4] >> 6) << 4);
    107
    0
        }
    108
    109
    0
        (d, dmin, scales, mins)
    110
    0
    }
    111
    112
    #[cfg(test)]
    113
    mod tests {
    114
        use super::gemv::compute_chunk_q4k_scalar;
    115
        use super::*;
    116
    117
        /// Golden Test: Fused kernel must match dequant+matmul within ε = 1e-3
    118
        /// This is the core falsification test from Section 12.4 of the spec.
    119
        #[test]
    120
        fn test_fused_q4k_golden_parity() {
    121
            // Create synthetic Q4K data (one super-block = 256 elements)
    122
            let in_dim = 256;
    123
            let out_dim = 4;
    124
            let num_blocks = 1;
    125
    126
            // Build Q4K test data
    127
            let mut q4k_data = Vec::with_capacity(out_dim * num_blocks * SUPER_BLOCK_BYTES);
    128
    129
            for row in 0..out_dim {
    130
                // d = 0.1, dmin = 0.05 (as f16)
    131
                let d: u16 = 0x2E66; // ~0.1 in f16
    132
                let dmin: u16 = 0x2A66; // ~0.05 in f16
    133
                q4k_data.extend_from_slice(&d.to_le_bytes());
    134
                q4k_data.extend_from_slice(&dmin.to_le_bytes());
    135
    136
                // Scales and mins (all set to 1 for simplicity)
    137
                let scales_packed = [0x01u8; 12];
    138
                q4k_data.extend_from_slice(&scales_packed);
    139
    140
                // Quantized values: pattern based on row
    141
                let mut qs = [0u8; 128];
    142
                for (i, q) in qs.iter_mut().enumerate() {
    143
                    // Low nibble: (row + i) % 16, High nibble: (row + i + 1) % 16
    144
                    let low = ((row + i) % 16) as u8;
    145
                    let high = ((row + i + 1) % 16) as u8;
    146
                    *q = low | (high << 4);
    147
                }
    148
                q4k_data.extend_from_slice(&qs);
    149
            }
    150
    151
            // Create input vector
    152
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.01).collect();
    153
    154
            // Compute using fused kernel
    155
            let fused_output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    156
    157
            // Compute reference: dequant then matmul
    158
            let mut reference_output = vec![0.0f32; out_dim];
    159
            for row in 0..out_dim {
    160
                let row_start = row * SUPER_BLOCK_BYTES;
    161
                let row_q4k = &q4k_data[row_start..row_start + SUPER_BLOCK_BYTES];
    162
                let f32_weights = dequantize_q4k_to_f32(row_q4k, in_dim);
    163
    164
                let mut sum = 0.0f32;
    165
                for (w, x) in f32_weights.iter().zip(input.iter()) {
    166
                    sum += w * x;
    167
                }
    168
                reference_output[row] = sum;
    169
            }
    170
    171
            // Golden parity check: |fused - reference| < 1e-3
    172
            for (i, (fused, reference)) in fused_output.iter().zip(reference_output.iter()).enumerate()
    173
            {
    174
                let diff = (fused - reference).abs();
    175
                assert!(
    176
                    diff < 1e-3,
    177
                    "Row {}: Fused kernel divergence: {} vs {} (Δ={})",
    178
                    i,
    179
                    fused,
    180
                    reference,
    181
                    diff
    182
                );
    183
            }
    184
        }
    185
    186
        /// Test scalar implementation matches optimized version
    187
        #[test]
    188
        fn test_scalar_vs_optimized_parity() {
    189
            let in_dim = 256;
    190
            let out_dim = 2;
    191
    192
            // Build simple Q4K test data
    193
            let mut q4k_data = Vec::new();
    194
            for _ in 0..out_dim {
    195
                q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    196
                q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin = 0
    197
                q4k_data.extend_from_slice(&[0x01u8; 12]); // scales
    198
                q4k_data.extend_from_slice(&[0x55u8; 128]); // qs = 5 | (5 << 4)
    199
            }
    200
    201
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.001).collect();
    202
    203
            let scalar_output = matmul_q4k_f32_scalar(&q4k_data, &input, out_dim, in_dim);
    204
            let optimized_output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    205
    206
            for (i, (s, o)) in scalar_output.iter().zip(optimized_output.iter()).enumerate() {
    207
                let diff = (s - o).abs();
    208
                // Allow small FP differences from mul_add vs separate multiply-add
    209
                assert!(
    210
                    diff < 1e-4,
    211
                    "Row {}: Scalar vs optimized divergence: {} vs {} (Δ={})",
    212
                    i,
    213
                    s,
    214
                    o,
    215
                    diff
    216
                );
    217
            }
    218
        }
    219
    220
        /// Test that output contains no NaN or Inf
    221
        #[test]
    222
        fn test_no_nan_inf() {
    223
            let in_dim = 256;
    224
            let out_dim = 4;
    225
    226
            let mut q4k_data = Vec::new();
    227
            for _ in 0..out_dim {
    228
                q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    229
                q4k_data.extend_from_slice(&[0x00, 0x38]); // dmin ~ 0.5
    230
                q4k_data.extend_from_slice(&[0x3Fu8; 12]); // max scales
    231
                q4k_data.extend_from_slice(&[0xFFu8; 128]); // max qs
    232
            }
    233
    234
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.01).collect();
    235
            let output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    236
    237
            for (i, &val) in output.iter().enumerate() {
    238
                assert!(val.is_finite(), "Row {}: Output is not finite: {}", i, val);
    239
            }
    240
        }
    241
    242
        /// Test AVX2 implementation matches scalar within tolerance
    243
        #[cfg(target_arch = "x86_64")]
    244
        #[test]
    245
        fn test_avx2_vs_scalar_parity() {
    246
            if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
    247
                eprintln!("Skipping AVX2 test - CPU doesn't support AVX2+FMA");
    248
                return;
    249
            }
    250
    251
            let in_dim = 512; // 2 super-blocks
    252
            let out_dim = 4;
    253
    254
            // Build Q4K test data with varied values
    255
            let mut q4k_data = Vec::new();
    256
            for row in 0..out_dim {
    257
                // d ~ 0.1, dmin ~ 0.05
    258
                q4k_data.extend_from_slice(&[0x66, 0x2E]); // d
    259
                q4k_data.extend_from_slice(&[0x66, 0x2A]); // dmin
    260
                // Varied scales
    261
                let scale_val = (row as u8 + 1) | ((row as u8 + 2) << 4);
    262
                q4k_data.extend_from_slice(&[scale_val; 12]);
    263
                // Varied quantized values
    264
                for i in 0..128 {
    265
                    let low = ((row + i) % 16) as u8;
    266
                    let high = ((row + i + 3) % 16) as u8;
    267
                    q4k_data.push(low | (high << 4));
    268
                }
    269
            }
    270
            // Duplicate for second super-block
    271
            let single_row_bytes = q4k_data.len() / out_dim;
    272
            let mut full_data = Vec::with_capacity(out_dim * single_row_bytes * 2);
    273
            for row in 0..out_dim {
    274
                let row_start = row * single_row_bytes;
    275
                full_data.extend_from_slice(&q4k_data[row_start..row_start + single_row_bytes]);
    276
                full_data.extend_from_slice(&q4k_data[row_start..row_start + single_row_bytes]);
    277
            }
    278
    279
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.002 - 0.5).collect();
    280
    281
            let scalar_output = matmul_q4k_f32(&full_data, &input, out_dim, in_dim);
    282
            let dispatch_output = matmul_q4k_f32_dispatch(&full_data, &input, out_dim, in_dim);
    283
    284
            for (i, (scalar, dispatch)) in scalar_output.iter().zip(dispatch_output.iter()).enumerate() {
    285
                let diff = (scalar - dispatch).abs();
    286
                let rel_diff = if scalar.abs() > 1e-6 {
    287
                    diff / scalar.abs()
    288
                } else {
    289
                    diff
    290
                };
    291
                // Allow 1e-5 relative error for FMA differences
    292
                assert!(
    293
                    rel_diff < 1e-5 || diff < 1e-5,
    294
                    "Row {}: AVX2 vs scalar divergence: {} vs {} (Δ={}, rel={})",
    295
                    i,
    296
                    dispatch,
    297
                    scalar,
    298
                    diff,
    299
                    rel_diff
    300
                );
    301
            }
    302
        }
    303
    304
        /// Test determinism: same input produces same output
    305
        #[test]
    306
        fn test_determinism() {
    307
            let in_dim = 256;
    308
            let out_dim = 2;
    309
    310
            let mut q4k_data = Vec::new();
    311
            for _ in 0..out_dim {
    312
                q4k_data.extend_from_slice(&[0x66, 0x2E]); // d
    313
                q4k_data.extend_from_slice(&[0x66, 0x2A]); // dmin
    314
                q4k_data.extend_from_slice(&[0x15u8; 12]);
    315
                q4k_data.extend_from_slice(&[0xABu8; 128]);
    316
            }
    317
    318
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.005).collect();
    319
    320
            let output1 = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    321
            let output2 = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    322
    323
            for (i, (a, b)) in output1.iter().zip(output2.iter()).enumerate() {
    324
                assert_eq!(
    325
                    a.to_bits(),
    326
                    b.to_bits(),
    327
                    "Row {}: Non-deterministic output: {} vs {}",
    328
                    i,
    329
                    a,
    330
                    b
    331
                );
    332
            }
    333
        }
    334
    335
        /// Test f16 conversion correctness
    336
        #[test]
    337
        fn test_f16_to_f32() {
    338
            // Test normal values
    339
            assert!((f16_to_f32(0x3C00) - 1.0).abs() < 1e-3); // 1.0
    340
            assert!((f16_to_f32(0x4000) - 2.0).abs() < 1e-3); // 2.0
    341
            assert!((f16_to_f32(0x3800) - 0.5).abs() < 1e-3); // 0.5
    342
    343
            // Test zero
    344
            assert_eq!(f16_to_f32(0x0000), 0.0);
    345
            assert_eq!(f16_to_f32(0x8000), -0.0);
    346
    347
            // Test subnormals (small values)
    348
            let small = f16_to_f32(0x0001);
    349
            assert!(small > 0.0 && small < 1e-4);
    350
        }
    351
    352
        #[test]
    353
        fn test_f16_to_f32_infinity_nan() {
    354
            // Positive infinity = 0x7C00
    355
            let inf = f16_to_f32(0x7C00);
    356
            assert!(inf.is_infinite() && inf.is_sign_positive());
    357
    358
            // Negative infinity = 0xFC00
    359
            let neg_inf = f16_to_f32(0xFC00);
    360
            assert!(neg_inf.is_infinite() && neg_inf.is_sign_negative());
    361
    362
            // Negative value
    363
            let neg_one = f16_to_f32(0xBC00); // -1.0
    364
            assert!((neg_one + 1.0).abs() < 1e-3);
    365
        }
    366
    367
        #[test]
    368
        fn test_dequantize_q4k_to_f32_basic() {
    369
            // Create a single Q4K block (144 bytes for 256 elements)
    370
            let mut block = vec![0u8; SUPER_BLOCK_BYTES];
    371
            // d = 1.0 (0x3C00)
    372
            block[0] = 0x00;
    373
            block[1] = 0x3C;
    374
            // dmin = 0 (0x0000)
    375
            block[2] = 0x00;
    376
            block[3] = 0x00;
    377
            // scales = all zeros
    378
            block[4..16].fill(0x00);
    379
            // qs = 0x55 (5 | 5<<4) for all values
    380
            block[16..144].fill(0x55);
    381
    382
            let result = dequantize_q4k_to_f32(&block, 256);
    383
            assert_eq!(result.len(), 256);
    384
    385
            // All values should be finite
    386
            for val in &result {
    387
                assert!(val.is_finite());
    388
            }
    389
        }
    390
    391
        #[test]
    392
        fn test_dequantize_q4k_to_f32_varies_scales() {
    393
            let mut block = vec![0u8; SUPER_BLOCK_BYTES];
    394
            block[0] = 0x00;
    395
            block[1] = 0x3C; // d = 1.0
    396
            block[2] = 0x00;
    397
            block[3] = 0x00; // dmin = 0
    398
    399
            // Set different scales for each group
    400
            for i in 0..12 {
    401
                block[4 + i] = (i * 10) as u8;
    402
            }
    403
    404
            // Set quantized values
    405
            block[16..144].fill(0x33); // 3 | 3<<4
    406
    407
            let result = dequantize_q4k_to_f32(&block, 256);
    408
            assert_eq!(result.len(), 256);
    409
            for val in &result {
    410
                assert!(val.is_finite());
    411
            }
    412
        }
    413
    414
        #[test]
    415
        fn test_matmul_q4k_f32_colmajor_basic() {
    416
            let in_dim = 256;
    417
            let out_dim = 2;
    418
    419
            let mut q4k_data = Vec::new();
    420
            for _ in 0..out_dim {
    421
                q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    422
                q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin = 0
    423
                q4k_data.extend_from_slice(&[0x01u8; 12]); // scales
    424
                q4k_data.extend_from_slice(&[0x55u8; 128]); // qs
    425
            }
    426
    427
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.01).collect();
    428
            let output = matmul_q4k_f32_colmajor(&q4k_data, &input, out_dim, in_dim);
    429
    430
            assert_eq!(output.len(), out_dim);
    431
            for val in &output {
    432
                assert!(val.is_finite());
    433
            }
    434
        }
    435
    436
        #[test]
    437
        fn test_matmul_q4k_f32_colmajor_dispatch_basic() {
    438
            let in_dim = 256;
    439
            let out_dim = 4;
    440
    441
            let mut q4k_data = Vec::new();
    442
            for row in 0..out_dim {
    443
                q4k_data.extend_from_slice(&[0x66, 0x2E]); // d ~ 0.1
    444
                q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin = 0
    445
                q4k_data.extend_from_slice(&[(row as u8 + 1); 12]); // varying scales
    446
                q4k_data.extend_from_slice(&[(row as u8 * 17).wrapping_add(0x44); 128]); // varying qs
    447
            }
    448
    449
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.01 - 1.0).collect();
    450
            let output = matmul_q4k_f32_colmajor_dispatch(&q4k_data, &input, out_dim, in_dim);
    451
    452
            assert_eq!(output.len(), out_dim);
    453
            for val in &output {
    454
                assert!(val.is_finite());
    455
            }
    456
        }
    457
    458
        #[test]
    459
        fn test_matmul_q4k_colmajor_produces_finite() {
    460
            // Column-major layout test: verify it produces valid finite outputs
    461
            // Note: colmajor and rowmajor have different data layout assumptions
    462
            let in_dim = 256;
    463
            let out_dim = 2;
    464
    465
            let mut q4k_data = Vec::new();
    466
            for _ in 0..out_dim {
    467
                q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    468
                q4k_data.extend_from_slice(&[0x00, 0x38]); // dmin ~ 0.5
    469
                q4k_data.extend_from_slice(&[0x01u8; 12]);
    470
                q4k_data.extend_from_slice(&[0x55u8; 128]);
    471
            }
    472
    473
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.005).collect();
    474
    475
            let rowmajor = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    476
            let colmajor = matmul_q4k_f32_colmajor(&q4k_data, &input, out_dim, in_dim);
    477
    478
            // Both should produce finite results
    479
            for (i, r) in rowmajor.iter().enumerate() {
    480
                assert!(r.is_finite(), "Row {}: rowmajor non-finite", i);
    481
            }
    482
            for (i, c) in colmajor.iter().enumerate() {
    483
                assert!(c.is_finite(), "Row {}: colmajor non-finite", i);
    484
            }
    485
        }
    486
    487
        #[test]
    488
        fn test_matmul_q4k_unaligned_dimensions() {
    489
            // Test with dimensions not aligned to block size (256)
    490
            let in_dim = 300;
    491
            let out_dim = 3;
    492
            let num_blocks = (in_dim + 255) / 256; // = 2 blocks
    493
    494
            let mut q4k_data = Vec::new();
    495
            for _ in 0..out_dim {
    496
                for _ in 0..num_blocks {
    497
                    q4k_data.extend_from_slice(&[0x00, 0x3C]); // d
    498
                    q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin
    499
                    q4k_data.extend_from_slice(&[0x01u8; 12]); // scales
    500
                    q4k_data.extend_from_slice(&[0x33u8; 128]); // qs
    501
                }
    502
            }
    503
    504
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.01).collect();
    505
            let output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    506
    507
            assert_eq!(output.len(), out_dim);
    508
            for val in &output {
    509
                assert!(val.is_finite());
    510
            }
    511
        }
    512
    513
        #[test]
    514
        fn test_matmul_q4k_zero_input() {
    515
            let in_dim = 256;
    516
            let out_dim = 2;
    517
    518
            let mut q4k_data = Vec::new();
    519
            for _ in 0..out_dim {
    520
                q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    521
                q4k_data.extend_from_slice(&[0x00, 0x38]); // dmin ~ 0.5
    522
                q4k_data.extend_from_slice(&[0x7Fu8; 12]); // max scales
    523
                q4k_data.extend_from_slice(&[0xFFu8; 128]); // max qs
    524
            }
    525
    526
            let input: Vec<f32> = vec![0.0; in_dim];
    527
            let output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    528
    529
            assert_eq!(output.len(), out_dim);
    530
            for val in &output {
    531
                assert_eq!(*val, 0.0, "Output should be zero when input is zero");
    532
            }
    533
        }
    534
    535
        #[test]
    536
        fn test_matmul_q4k_large_dimensions() {
    537
            let in_dim = 1024;
    538
            let out_dim = 8;
    539
            let num_blocks = in_dim / 256;
    540
    541
            let mut q4k_data = Vec::new();
    542
            for row in 0..out_dim {
    543
                for blk in 0..num_blocks {
    544
                    let val = ((row * num_blocks + blk) as u8).wrapping_mul(17);
    545
                    q4k_data.extend_from_slice(&[0x66, 0x2E]); // d ~ 0.1
    546
                    q4k_data.extend_from_slice(&[0x33, 0x2A]); // dmin ~ 0.05
    547
                    q4k_data.extend_from_slice(&[(val.wrapping_add(1)); 12]);
    548
                    q4k_data.extend_from_slice(&[val.wrapping_add(0x55); 128]);
    549
                }
    550
            }
    551
    552
            let input: Vec<f32> = (0..in_dim).map(|i| ((i % 100) as f32) * 0.01).collect();
    553
            let output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    554
    555
            assert_eq!(output.len(), out_dim);
    556
            for val in &output {
    557
                assert!(val.is_finite());
    558
            }
    559
        }
    560
    561
        #[test]
    562
        fn test_parse_q4k_header() {
    563
            let mut block = vec![0u8; 144];
    564
            // d = 1.0 (0x3C00), dmin = 0.5 (0x3800)
    565
            block[0] = 0x00;
    566
            block[1] = 0x3C;
    567
            block[2] = 0x00;
    568
            block[3] = 0x38;
    569
            // scales_bytes[0..12] for llama.cpp format
    570
            // bytes 0-3: lower 6 bits = scales[0-3], upper 2 bits = scales[4-7] upper bits
    571
            // bytes 4-7: lower 6 bits = mins[0-3], upper 2 bits = mins[4-7] upper bits
    572
            // bytes 8-11: lower 4 bits = scales[4-7] lower bits, upper 4 bits = mins[4-7] lower bits
    573
            block[4..8].copy_from_slice(&[0x01, 0x02, 0x03, 0x04]); // scales[0-3] = 1,2,3,4
    574
            block[8..12].copy_from_slice(&[0x0A, 0x0B, 0x0C, 0x0D]); // mins[0-3] = 10,11,12,13
    575
            block[12..16].copy_from_slice(&[0x55, 0x66, 0x77, 0x88]); // combined lower nibbles
    576
    577
            let (d, dmin, scales, mins) = parse_q4k_header(&block);
    578
    579
            assert!((d - 1.0).abs() < 0.01, "d should be ~1.0, got {}", d);
    580
            assert!((dmin - 0.5).abs() < 0.01, "dmin should be ~0.5, got {}", dmin);
    581
            // Check first scales/mins have expected low 6-bit values
    582
            assert_eq!(scales[0], 0x01, "scales[0] should be 1");
    583
            assert_eq!(scales[1], 0x02, "scales[1] should be 2");
    584
            assert_eq!(mins[0], 0x0A, "mins[0] should be 10");
    585
            assert_eq!(mins[1], 0x0B, "mins[1] should be 11");
    586
        }
    587
    588
        #[test]
    589
        fn test_matmul_q4k_single_row() {
    590
            let in_dim = 256;
    591
            let out_dim = 1;
    592
    593
            let mut q4k_data = Vec::new();
    594
            q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    595
            q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin = 0
    596
            q4k_data.extend_from_slice(&[0x01u8; 12]); // scales
    597
            q4k_data.extend_from_slice(&[0xAAu8; 128]); // qs
    598
    599
            let input: Vec<f32> = vec![1.0; in_dim];
    600
            let output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    601
    602
            assert_eq!(output.len(), 1);
    603
            assert!(output[0].is_finite());
    604
        }
    605
    606
        /// Test AVX2 matmul with large dimensions (exercises full SIMD paths)
    607
        #[cfg(target_arch = "x86_64")]
    608
        #[test]
    609
        fn test_avx2_large_matrix_mul() {
    610
            if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
    611
                eprintln!("Skipping AVX2 large matrix test - CPU doesn't support AVX2+FMA");
    612
                return;
    613
            }
    614
    615
            let in_dim = 4096; // 16 super-blocks
    616
            let out_dim = 32;
    617
    618
            // Build Q4K test data with realistic values
    619
            let mut q4k_data = Vec::new();
    620
            for row in 0..out_dim {
    621
                for _sb in 0..(in_dim / 256) {
    622
                    // d ~ 0.1, dmin ~ 0.05
    623
                    q4k_data.extend_from_slice(&[0x66, 0x2E]); // d
    624
                    q4k_data.extend_from_slice(&[0x66, 0x2A]); // dmin
    625
                    // Varied scales based on row
    626
                    let scale_val = (row as u8 % 16) | (((row + 1) as u8 % 16) << 4);
    627
                    q4k_data.extend_from_slice(&[scale_val; 12]);
    628
                    // Varied quantized values
    629
                    for i in 0..128 {
    630
                        let low = ((row + i) % 16) as u8;
    631
                        let high = ((row + i + 3) % 16) as u8;
    632
                        q4k_data.push(low | (high << 4));
    633
                    }
    634
                }
    635
            }
    636
    637
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.001 - 2.0).collect();
    638
    639
            let scalar_output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    640
            let dispatch_output = matmul_q4k_f32_dispatch(&q4k_data, &input, out_dim, in_dim);
    641
    642
            for (i, (scalar, dispatch)) in scalar_output.iter().zip(dispatch_output.iter()).enumerate() {
    643
                let diff = (scalar - dispatch).abs();
    644
                let rel_diff = if scalar.abs() > 1e-6 {
    645
                    diff / scalar.abs()
    646
                } else {
    647
                    diff
    648
                };
    649
                assert!(
    650
                    rel_diff < 1e-4 || diff < 1e-4,
    651
                    "Row {}: AVX2 vs scalar divergence: {} vs {} (Δ={}, rel={})",
    652
                    i, dispatch, scalar, diff, rel_diff
    653
                );
    654
            }
    655
        }
    656
    657
        /// Test colmajor AVX2 path with realistic dimensions
    658
        #[cfg(target_arch = "x86_64")]
    659
        #[test]
    660
        fn test_avx2_colmajor_large() {
    661
            if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
    662
                eprintln!("Skipping AVX2 colmajor test - CPU doesn't support AVX2+FMA");
    663
                return;
    664
            }
    665
    666
            let in_dim = 2048; // 8 super-blocks
    667
            let out_dim = 16;
    668
    669
            let mut q4k_data = Vec::new();
    670
            for row in 0..out_dim {
    671
                for sb in 0..(in_dim / 256) {
    672
                    q4k_data.extend_from_slice(&[0x66, 0x2E]); // d
    673
                    q4k_data.extend_from_slice(&[0x33, 0x2A]); // dmin
    674
                    let scale_val = ((row + sb) as u8 % 16) | (((row + sb + 1) as u8 % 16) << 4);
    675
                    q4k_data.extend_from_slice(&[scale_val; 12]);
    676
                    for i in 0..128 {
    677
                        q4k_data.push(((i % 16) | ((i + 1) % 16 << 4)) as u8);
    678
                    }
    679
                }
    680
            }
    681
    682
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.002 - 1.0).collect();
    683
    684
            let output = matmul_q4k_f32_colmajor(&q4k_data, &input, out_dim, in_dim);
    685
            let output_dispatch = matmul_q4k_f32_colmajor_dispatch(&q4k_data, &input, out_dim, in_dim);
    686
    687
            assert_eq!(output.len(), out_dim);
    688
            assert_eq!(output_dispatch.len(), out_dim);
    689
    690
            for (i, (base, dispatched)) in output.iter().zip(output_dispatch.iter()).enumerate() {
    691
                let diff = (base - dispatched).abs();
    692
                assert!(
    693
                    diff < 1e-3 || (diff / base.abs()) < 1e-4,
    694
                    "Row {}: colmajor mismatch: {} vs {} (diff={})",
    695
                    i, base, dispatched, diff
    696
                );
    697
            }
    698
        }
    699
    700
        /// Test non-aligned dimensions (exercises scalar remainder handling)
    701
        #[cfg(target_arch = "x86_64")]
    702
        #[test]
    703
        fn test_avx2_non_aligned_dimensions() {
    704
            if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
    705
                eprintln!("Skipping AVX2 non-aligned test - CPU doesn't support AVX2+FMA");
    706
                return;
    707
            }
    708
    709
            // Non-aligned: 768 = 3 super-blocks (not power of 2)
    710
            let in_dim = 768;
    711
            let out_dim = 7; // Odd number
    712
    713
            let mut q4k_data = Vec::new();
    714
            for row in 0..out_dim {
    715
                for _sb in 0..(in_dim / 256) {
    716
                    q4k_data.extend_from_slice(&[0x66, 0x2E]);
    717
                    q4k_data.extend_from_slice(&[0x66, 0x2A]);
    718
                    let scale_val = (row as u8 % 16) | (((row + 1) as u8 % 16) << 4);
    719
                    q4k_data.extend_from_slice(&[scale_val; 12]);
    720
                    for i in 0..128 {
    721
                        q4k_data.push(((i % 16) | ((i + 5) % 16 << 4)) as u8);
    722
                    }
    723
                }
    724
            }
    725
    726
            let input: Vec<f32> = (0..in_dim).map(|i| ((i as f32) * 0.003).sin()).collect();
    727
    728
            let scalar_output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    729
            let dispatch_output = matmul_q4k_f32_dispatch(&q4k_data, &input, out_dim, in_dim);
    730
    731
            assert_eq!(scalar_output.len(), out_dim);
    732
            assert_eq!(dispatch_output.len(), out_dim);
    733
    734
            for (i, (scalar, dispatch)) in scalar_output.iter().zip(dispatch_output.iter()).enumerate() {
    735
                let diff = (scalar - dispatch).abs();
    736
                let rel_diff = if scalar.abs() > 1e-6 {
    737
                    diff / scalar.abs()
    738
                } else {
    739
                    diff
    740
                };
    741
                // FMA operations can have ordering differences, allow 1e-5 relative error
    742
                assert!(
    743
                    rel_diff < 1e-5 || diff < 1e-2,
    744
                    "Row {}: non-aligned AVX2 mismatch: {} vs {} (diff={}, rel={})",
    745
                    i, scalar, dispatch, diff, rel_diff
    746
                );
    747
            }
    748
        }
    749
    750
        /// Test parallel SIMD execution (exercises compute_chunk_q4k_avx2)
    751
        #[cfg(all(target_arch = "x86_64", feature = "parallel"))]
    752
        #[test]
    753
        fn test_parallel_avx2_large_batch() {
    754
            if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
    755
                eprintln!("Skipping parallel AVX2 test - CPU doesn't support AVX2+FMA");
    756
                return;
    757
            }
    758
    759
            // Large enough to trigger parallel path (>1000 rows)
    760
            let in_dim = 1024;
    761
            let out_dim = 2048; // Large output dim for parallel execution
    762
    763
            let mut q4k_data = Vec::new();
    764
            for row in 0..out_dim {
    765
                for _sb in 0..(in_dim / 256) {
    766
                    q4k_data.extend_from_slice(&[0x66, 0x2E]);
    767
                    q4k_data.extend_from_slice(&[0x33, 0x2A]);
    768
                    let scale_val = ((row % 256) as u8) | (((row / 256) % 16) as u8 * 16);
    769
                    q4k_data.extend_from_slice(&[scale_val; 12]);
    770
                    for i in 0..128 {
    771
                        q4k_data.push(((i * row) % 256) as u8);
    772
                    }
    773
                }
    774
            }
    775
    776
            let input: Vec<f32> = (0..in_dim).map(|i| (i as f32) * 0.001).collect();
    777
    778
            let output = matmul_q4k_f32_colmajor_dispatch(&q4k_data, &input, out_dim, in_dim);
    779
    780
            assert_eq!(output.len(), out_dim);
    781
            for (i, val) in output.iter().enumerate() {
    782
                assert!(
    783
                    val.is_finite(),
    784
                    "Row {}: parallel AVX2 produced non-finite: {}",
    785
                    i, val
    786
                );
    787
            }
    788
        }
    789
    790
        // =========================================================================
    791
        // Golden Vector Tests (Section 12.4: Q4K fused matmul ≈ dequant+f32_matmul)
    792
        // =========================================================================
    793
    794
        /// Helper: naive f32 matrix-vector multiplication
    795
        fn matmul_f32_naive(weights: &[f32], input: &[f32], out_dim: usize, in_dim: usize) -> Vec<f32> {
    796
            let mut output = vec![0.0f32; out_dim];
    797
            for row in 0..out_dim {
    798
                let mut sum = 0.0f32;
    799
                for col in 0..in_dim {
    800
                    sum += weights[row * in_dim + col] * input[col];
    801
                }
    802
                output[row] = sum;
    803
            }
    804
            output
    805
        }
    806
    807
        /// Golden Vector Test: Q4K matmul ≈ dequant + f32 matmul
    808
        ///
    809
        /// This test verifies the invariant from Section 12.4:
    810
        /// matmul_q4k_f32(W, x) ≈ matmul(dequant_q4k_to_f32(W), x) within ε
    811
        ///
    812
        /// Quantization introduces error, so we use a relaxed tolerance (5%).
    813
        #[test]
    814
        fn test_golden_vector_q4k_matmul_vs_dequant() {
    815
            use crate::backends::q4k::dequantize_q4k_to_f32;
    816
    817
            // Realistic dimensions for LLM layers
    818
            let in_dim = 512; // 2 super-blocks
    819
            let out_dim = 8;
    820
    821
            // Build Q4K test data with realistic distribution
    822
            let mut q4k_data = Vec::new();
    823
            for row in 0..out_dim {
    824
                for sb in 0..(in_dim / 256) {
    825
                    // d ~ 0.1, dmin ~ 0.05 (realistic for normalized weights)
    826
                    q4k_data.extend_from_slice(&[0x66, 0x2E]); // d
    827
                    q4k_data.extend_from_slice(&[0x66, 0x2A]); // dmin
    828
                    // Varied scales based on position
    829
                    let scale_base = ((row * 7 + sb * 3) % 16) as u8;
    830
                    for i in 0..12 {
    831
                        q4k_data.push(scale_base + (i as u8 % 4));
    832
                    }
    833
                    // Varied quantized values (4-bit, so 0-15)
    834
                    for i in 0..128 {
    835
                        let low = ((row + sb + i) % 16) as u8;
    836
                        let high = ((row + sb + i + 5) % 16) as u8;
    837
                        q4k_data.push(low | (high << 4));
    838
                    }
    839
                }
    840
            }
    841
    842
            // Random-ish input vector (sinusoidal distribution)
    843
            let input: Vec<f32> = (0..in_dim)
    844
                .map(|i| ((i as f32) * 0.017).sin() * 0.5)
    845
                .collect();
    846
    847
            // Method 1: Fused Q4K matmul
    848
            let fused_output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    849
    850
            // Method 2: Dequantize + f32 matmul
    851
            let total_elements = in_dim * out_dim;
    852
            let dequantized_weights = dequantize_q4k_to_f32(&q4k_data, total_elements);
    853
            let reference_output = matmul_f32_naive(&dequantized_weights, &input, out_dim, in_dim);
    854
    855
            // Verify Golden Invariant: error within 5% or absolute 0.01
    856
            assert_eq!(fused_output.len(), reference_output.len());
    857
            let mut max_rel_error = 0.0f32;
    858
            let mut max_abs_error = 0.0f32;
    859
    860
            for (i, (fused, reference)) in fused_output.iter().zip(reference_output.iter()).enumerate() {
    861
                let abs_error = (fused - reference).abs();
    862
                let rel_error = if reference.abs() > 1e-6 {
    863
                    abs_error / reference.abs()
    864
                } else {
    865
                    abs_error
    866
                };
    867
                max_rel_error = max_rel_error.max(rel_error);
    868
                max_abs_error = max_abs_error.max(abs_error);
    869
    870
                assert!(
    871
                    rel_error < 0.05 || abs_error < 0.01,
    872
                    "Golden invariant violated at row {}: fused={}, reference={}, \
    873
                     rel_error={:.4}%, abs_error={:.6}",
    874
                    i, fused, reference, rel_error * 100.0, abs_error
    875
                );
    876
            }
    877
    878
            // Report max errors for visibility
    879
            eprintln!(
    880
                "[Golden Q4K Test] max_rel_error={:.4}%, max_abs_error={:.6}",
    881
                max_rel_error * 100.0, max_abs_error
    882
            );
    883
        }
    884
    885
        /// Golden Vector Test: dispatch path also satisfies invariant
    886
        #[test]
    887
        fn test_golden_vector_q4k_dispatch_vs_dequant() {
    888
            use crate::backends::q4k::dequantize_q4k_to_f32;
    889
    890
            // Larger dimensions to exercise SIMD paths
    891
            let in_dim = 1024; // 4 super-blocks
    892
            let out_dim = 16;
    893
    894
            let mut q4k_data = Vec::new();
    895
            for row in 0..out_dim {
    896
                for sb in 0..(in_dim / 256) {
    897
                    q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    898
                    q4k_data.extend_from_slice(&[0x00, 0x38]); // dmin ~ 0.5
    899
                    for i in 0..12 {
    900
                        q4k_data.push(((row + sb + i) % 64) as u8);
    901
                    }
    902
                    for i in 0..128 {
    903
                        let low = ((row * 3 + sb * 7 + i) % 16) as u8;
    904
                        let high = ((row * 5 + sb * 11 + i * 2) % 16) as u8;
    905
                        q4k_data.push(low | (high << 4));
    906
                    }
    907
                }
    908
            }
    909
    910
            let input: Vec<f32> = (0..in_dim)
    911
                .map(|i| ((i as f32) * 0.013 + 0.5).cos() * 0.3)
    912
                .collect();
    913
    914
            // Dispatch (may use AVX2/SIMD)
    915
            let dispatch_output = matmul_q4k_f32_dispatch(&q4k_data, &input, out_dim, in_dim);
    916
    917
            // Reference: dequantize + f32
    918
            let total_elements = in_dim * out_dim;
    919
            let dequantized = dequantize_q4k_to_f32(&q4k_data, total_elements);
    920
            let reference_output = matmul_f32_naive(&dequantized, &input, out_dim, in_dim);
    921
    922
            let mut max_rel_error = 0.0f32;
    923
            for (i, (dispatch, reference)) in dispatch_output.iter().zip(reference_output.iter()).enumerate() {
    924
                let abs_error = (dispatch - reference).abs();
    925
                let rel_error = if reference.abs() > 1e-6 {
    926
                    abs_error / reference.abs()
    927
                } else {
    928
                    abs_error
    929
                };
    930
                max_rel_error = max_rel_error.max(rel_error);
    931
    932
                assert!(
    933
                    rel_error < 0.05 || abs_error < 0.01,
    934
                    "Golden invariant violated (dispatch) at row {}: \
    935
                     dispatch={}, reference={}, rel_error={:.4}%",
    936
                    i, dispatch, reference, rel_error * 100.0
    937
                );
    938
            }
    939
    940
            eprintln!(
    941
                "[Golden Q4K Dispatch Test] max_rel_error={:.4}%",
    942
                max_rel_error * 100.0
    943
            );
    944
        }
    945
    946
        /// Edge case: zero input vector should produce zero output
    947
        #[test]
    948
        fn test_golden_vector_zero_input() {
    949
            let in_dim = 256;
    950
            let out_dim = 4;
    951
    952
            let mut q4k_data = Vec::new();
    953
            for _row in 0..out_dim {
    954
                q4k_data.extend_from_slice(&[0x66, 0x2E]); // d
    955
                q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin = 0 (important for zero output)
    956
                q4k_data.extend_from_slice(&[0x01u8; 12]);
    957
                q4k_data.extend_from_slice(&[0x55u8; 128]); // Non-zero weights
    958
            }
    959
    960
            let input = vec![0.0f32; in_dim];
    961
            let output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    962
    963
            // With dmin=0 and all-zero input, output should be near zero
    964
            for (i, val) in output.iter().enumerate() {
    965
                assert!(
    966
                    val.abs() < 1e-6,
    967
                    "Zero input should give ~zero output, got {} at row {}",
    968
                    val, i
    969
                );
    970
            }
    971
        }
    972
    973
        /// Edge case: uniform input vector
    974
        #[test]
    975
        fn test_golden_vector_uniform_input() {
    976
            use crate::backends::q4k::dequantize_q4k_to_f32;
    977
    978
            let in_dim = 256;
    979
            let out_dim = 2;
    980
    981
            let mut q4k_data = Vec::new();
    982
            for row in 0..out_dim {
    983
                q4k_data.extend_from_slice(&[0x00, 0x3C]); // d ~ 1.0
    984
                q4k_data.extend_from_slice(&[0x00, 0x00]); // dmin = 0
    985
                q4k_data.extend_from_slice(&[0x01u8; 12]);
    986
                // Uniform quantized weights
    987
                q4k_data.extend_from_slice(&[((row + 1) * 0x11) as u8; 128]);
    988
            }
    989
    990
            let input = vec![1.0f32; in_dim]; // All ones
    991
            let fused_output = matmul_q4k_f32(&q4k_data, &input, out_dim, in_dim);
    992
    993
            let total_elements = in_dim * out_dim;
    994
            let dequantized = dequantize_q4k_to_f32(&q4k_data, total_elements);
    995
            let reference_output = matmul_f32_naive(&dequantized, &input, out_dim, in_dim);
    996
    997
            for (i, (fused, reference)) in fused_output.iter().zip(reference_output.iter()).enumerate() {
    998
                let rel_error = if reference.abs() > 1e-6 {
    999
                    (fused - reference).abs() / reference.abs()
    1000
                } else {
    1001
                    (fused - reference).abs()
    1002
                };
    1003
                assert!(
    1004
                    rel_error < 0.05,
    1005
                    "Uniform input failed at row {}: fused={}, ref={}, err={:.2}%",
    1006
                    i, fused, reference, rel_error * 100.0
    1007
                );
    1008
            }
    1009
        }
    1010
    1011
        #[test]
    1012
        fn test_parallel_dispatch_large_matrix() {
    1013
            // Test parallel path: total_work >= 8_000_000
    1014
            // Use 4096 x 2048 = 8_388_608 ops (triggers parallel)
    1015
            let out_dim = 4096;
    1016
            let in_dim = 2048; // Must be multiple of 256 (SUPER_BLOCK_SIZE)
    1017
            let total_work = out_dim * in_dim;
    1018
            assert!(
    1019
                total_work >= 8_000_000,
    1020
                "Test must trigger parallel path"
    1021
            );
    1022
    1023
            let num_superblocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    1024
            let row_bytes = num_superblocks_per_row * SUPER_BLOCK_BYTES;
    1025
            let total_bytes = out_dim * row_bytes;
    1026
    1027
            // Create deterministic test data
    1028
            let mut q4k_data = vec![0u8; total_bytes];
    1029
            for row in 0..out_dim {
    1030
                for sb in 0..num_superblocks_per_row {
    1031
                    let offset = row * row_bytes + sb * SUPER_BLOCK_BYTES;
    1032
                    // d = 1.0 as f16
    1033
                    q4k_data[offset] = 0x00;
    1034
                    q4k_data[offset + 1] = 0x3C;
    1035
                    // dmin = 0.0
    1036
                    q4k_data[offset + 2] = 0x00;
    1037
                    q4k_data[offset + 3] = 0x00;
    1038
                    // scales = 1 for all
    1039
                    for i in 0..12 {
    1040
                        q4k_data[offset + 4 + i] = 0x01;
    1041
                    }
    1042
                    // qs = predictable pattern
    1043
                    for i in 0..128 {
    1044
                        q4k_data[offset + 16 + i] = ((row + sb + i) % 16) as u8;
    1045
                    }
    1046
                }
    1047
            }
    1048
    1049
            let input: Vec<f32> = (0..in_dim).map(|i| (i % 10) as f32 * 0.1).collect();
    1050
    1051
            // Call dispatch - should use parallel path
    1052
            let result = matmul_q4k_f32_dispatch(&q4k_data, &input, out_dim, in_dim);
    1053
    1054
            // Verify dimensions and finiteness
    1055
            assert_eq!(result.len(), out_dim);
    1056
            for (i, &val) in result.iter().enumerate() {
    1057
                assert!(
    1058
                    val.is_finite(),
    1059
                    "Result[{}] is not finite: {}",
    1060
                    i,
    1061
                    val
    1062
                );
    1063
            }
    1064
    1065
            // Compare a few rows against scalar for consistency
    1066
            let scalar_result = matmul_q4k_f32_scalar(&q4k_data, &input, out_dim, in_dim);
    1067
            for i in (0..out_dim).step_by(512) {
    1068
                let diff = (result[i] - scalar_result[i]).abs();
    1069
                let tol = scalar_result[i].abs() * 0.01 + 1e-5;
    1070
                assert!(
    1071
                    diff < tol,
    1072
                    "Parallel vs scalar mismatch at row {}: parallel={}, scalar={}, diff={}",
    1073
                    i,
    1074
                    result[i],
    1075
                    scalar_result[i],
    1076
                    diff
    1077
                );
    1078
            }
    1079
        }
    1080
    1081
        #[test]
    1082
        fn test_parallel_colmajor_large_matrix() {
    1083
            // Test colmajor path
    1084
            // ne0 = output dimension (rows), ne1 = input dimension (columns)
    1085
            // Input must have length ne1
    1086
            let ne0 = 2048; // output dimension (rows), must be multiple of 256
    1087
            let ne1 = 4096; // input dimension (columns)
    1088
    1089
            let blocks_per_col = (ne0 + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    1090
            let col_bytes = blocks_per_col * SUPER_BLOCK_BYTES;
    1091
            let total_bytes = ne1 * col_bytes;
    1092
    1093
            let mut q4k_data = vec![0u8; total_bytes];
    1094
            for col in 0..ne1 {
    1095
                for sb in 0..blocks_per_col {
    1096
                    let offset = col * col_bytes + sb * SUPER_BLOCK_BYTES;
    1097
                    // d = 0.5 as f16
    1098
                    q4k_data[offset] = 0x00;
    1099
                    q4k_data[offset + 1] = 0x38;
    1100
                    // dmin = 0.0
    1101
                    q4k_data[offset + 2] = 0x00;
    1102
                    q4k_data[offset + 3] = 0x00;
    1103
                    // scales
    1104
                    for i in 0..12 {
    1105
                        q4k_data[offset + 4 + i] = 0x02;
    1106
                    }
    1107
                    // qs
    1108
                    for i in 0..128 {
    1109
                        q4k_data[offset + 16 + i] = ((col ^ sb ^ i) % 16) as u8;
    1110
                    }
    1111
                }
    1112
            }
    1113
    1114
            // Input must have length ne1 (input dimension)
    1115
            let input: Vec<f32> = (0..ne1).map(|i| ((i % 7) as f32 - 3.0) * 0.1).collect();
    1116
    1117
            // Use colmajor dispatch
    1118
            let result = matmul_q4k_f32_colmajor_dispatch(&q4k_data, &input, ne0, ne1);
    1119
    1120
            // Output has ne0 elements
    1121
            assert_eq!(result.len(), ne0);
    1122
            for (i, &val) in result.iter().enumerate() {
    1123
                assert!(val.is_finite(), "Result[{}] is not finite: {}", i, val);
    1124
            }
    1125
        }
    1126
    1127
        #[test]
    1128
        fn test_compute_chunk_scalar_small() {
    1129
            // Directly test compute_chunk_q4k_scalar
    1130
            let in_dim = 256;
    1131
            let out_dim = 4;
    1132
            let num_blocks_per_row = 1;
    1133
            let row_bytes = SUPER_BLOCK_BYTES;
    1134
    1135
            let mut q4k_data = vec![0u8; out_dim * row_bytes];
    1136
            for row in 0..out_dim {
    1137
                let offset = row * row_bytes;
    1138
                // d = 1.0 as f16
    1139
                q4k_data[offset] = 0x00;
    1140
                q4k_data[offset + 1] = 0x3C;
    1141
                // dmin = 0.0
    1142
                q4k_data[offset + 2] = 0x00;
    1143
                q4k_data[offset + 3] = 0x00;
    1144
                // scales = 1
    1145
                for i in 0..12 {
    1146
                    q4k_data[offset + 4 + i] = 0x01;
    1147
                }
    1148
                // qs = all zeros (simplest case)
    1149
                for i in 0..128 {
    1150
                    q4k_data[offset + 16 + i] = 0x00;
    1151
                }
    1152
            }
    1153
    1154
            let input = vec![1.0f32; in_dim];
    1155
            let mut chunk = vec![0.0f32; out_dim];
    1156
    1157
            compute_chunk_q4k_scalar(
    1158
                &q4k_data,
    1159
                &input,
    1160
                &mut chunk,
    1161
                0,
    1162
                out_dim,
    1163
                in_dim,
    1164
                num_blocks_per_row,
    1165
                row_bytes,
    1166
            );
    1167
    1168
            // With qs=0, d=1, scales=1, dmin=0, result should be negative
    1169
            // Each element: d * scale * 0 - dmin * min = 0 - 0 = 0
    1170
            // Actually with all zeros in qs and dmin=0, output should be 0
    1171
            for (i, &val) in chunk.iter().enumerate() {
    1172
                assert!(
    1173
                    val.is_finite(),
    1174
                    "Chunk[{}] is not finite: {}",
    1175
                    i,
    1176
                    val
    1177
                );
    1178
            }
    1179
        }
    1180
    }