feat(quant/cpu): add NEON fused dequant+dot kernels for aarch64

Implement fused dequantization and dot product kernels for Q4_K,
Q5_K, and Q6_K formats using ARM NEON intrinsics. Adds a shared
horizontal sum helper (dot_f32.rs) used across all three kernels.

These kernels are the aarch64 counterpart to the existing x86 AVX2
implementations and are wired into the dispatch logic added in the
previous commit.

Author: farhan-syah

src/quant/cpu/kernels/simd/aarch64/dot_f32.rs

@@ -0,0 +1,48 @@
+//! NEON f32 dot product with FMA for quantized matmul accumulation
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+const F32_LANES: usize = 4;
+
+/// Horizontal sum of 4 f32 values in a NEON register.
+///
+/// # Safety
+/// - CPU must support NEON (always true on AArch64)
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+#[inline]
+pub unsafe fn hsum_f32_neon(v: float32x4_t) -> f32 {
+    let pair = vadd_f32(vget_low_f32(v), vget_high_f32(v));
+    vget_lane_f32::<0>(vpadd_f32(pair, pair))
+}
+
+/// NEON f32 dot product of two slices using FMA.
+///
+/// # Safety
+/// - CPU must support NEON (always true on AArch64)
+/// - `a` and `b` must be valid for reads of `len` f32 elements each
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+pub unsafe fn dot_f32_neon(a: *const f32, b: *const f32, len: usize) -> f32 {
+    let chunks = len / F32_LANES;
+    let remainder = len % F32_LANES;
+
+    let mut acc = vdupq_n_f32(0.0);
+
+    for i in 0..chunks {
+        let offset = i * F32_LANES;
+        let va = vld1q_f32(a.add(offset));
+        let vb = vld1q_f32(b.add(offset));
+        acc = vfmaq_f32(acc, va, vb);
+    }
+
+    let mut result = hsum_f32_neon(acc);
+
+    for i in 0..remainder {
+        let offset = chunks * F32_LANES + i;
+        result += *a.add(offset) * *b.add(offset);
+    }
+
+    result
+}

src/quant/cpu/kernels/simd/aarch64/fused_q4k.rs

@@ -0,0 +1,99 @@
+//! NEON fused dequant+dot for Q4_K format
+//!
+//! Q4_K block layout (256 elements, 144 bytes):
+//!   [0..2]   d (f16 scale)
+//!   [2..4]   dmin (f16 minimum)
+//!   [4..16]  sc (12-byte packed 6-bit scales+mins for 8 sub-blocks)
+//!   [16..144] qs (128 bytes of 4-bit quantized values, 2 per byte)
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+use half::f16;
+
+use super::super::super::dequant_k_quants::unpack_q4k_q5k_scales;
+use super::dot_f32::hsum_f32_neon;
+
+const F32_LANES: usize = 4;
+
+/// Fused dequant+dot for Q4_K using NEON.
+///
+/// # Safety
+/// Requires NEON. Caller must ensure act.len() >= k and blocks covers k/256 blocks.
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+pub unsafe fn fused_dot_q4k_neon(act: &[f32], blocks: &[u8], k: usize) -> f32 {
+    const BLOCK_SIZE: usize = 256;
+    const BLOCK_BYTES: usize = 144;
+    let num_blocks = k / BLOCK_SIZE;
+
+    debug_assert!(act.len() >= k, "act.len() {} < k {}", act.len(), k);
+    debug_assert!(
+        blocks.len() >= num_blocks * BLOCK_BYTES,
+        "blocks.len() {} < required {}",
+        blocks.len(),
+        num_blocks * BLOCK_BYTES
+    );
+
+    let mut total_acc = vdupq_n_f32(0.0);
+
+    for b in 0..num_blocks {
+        let block = &blocks[b * BLOCK_BYTES..];
+        let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
+        let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
+        let sc = &block[4..16];
+        let qs = &block[16..144];
+        let act_block = &act[b * BLOCK_SIZE..];
+
+        let (scales, mins) = unpack_q4k_q5k_scales(sc);
+
+        for j in 0..8 {
+            let dl = d * scales[j] as f32;
+            let ml = dmin * mins[j] as f32;
+
+            let chunk = j / 2;
+            let is_high = j % 2 == 1;
+            let qs_base = chunk * 32;
+
+            let act_sub = &act_block[j * 32..];
+            let dl_vec = vdupq_n_f32(dl);
+            let ml_vec = vdupq_n_f32(ml);
+            let mask_0f = vdupq_n_u32(0x0F);
+
+            // Process 32 elements in 8 groups of 4
+            for g in 0..8 {
+                let l_base = g * 4;
+
+                // Load 4 bytes of quantized data, zero-extend to u32
+                let q0 = qs[qs_base + l_base] as u32;
+                let q1 = qs[qs_base + l_base + 1] as u32;
+                let q2 = qs[qs_base + l_base + 2] as u32;
+                let q3 = qs[qs_base + l_base + 3] as u32;
+
+                let raw = vcreate_u32(q0 as u64 | (q1 as u64) << 32);
+                let raw_hi = vcreate_u32(q2 as u64 | (q3 as u64) << 32);
+                let raw256 = vcombine_u32(raw, raw_hi);
+
+                // Extract nibbles
+                let nibbles = if is_high {
+                    vandq_u32(vshrq_n_u32::<4>(raw256), mask_0f)
+                } else {
+                    vandq_u32(raw256, mask_0f)
+                };
+
+                // Convert to f32
+                let q_f32 = vcvtq_f32_u32(nibbles);
+
+                // Load 4 activation values
+                let a = vld1q_f32(act_sub.as_ptr().add(l_base));
+
+                // Accumulate: dl * (a * q) - ml * a
+                let aq = vmulq_f32(a, q_f32);
+                total_acc = vfmaq_f32(total_acc, dl_vec, aq);
+                total_acc = vsubq_f32(total_acc, vmulq_f32(ml_vec, a));
+            }
+        }
+    }
+
+    hsum_f32_neon(total_acc)
+}

src/quant/cpu/kernels/simd/aarch64/fused_q5k.rs

@@ -0,0 +1,95 @@
+//! NEON fused dequant+dot for Q5_K format
+//!
+//! Q5_K block layout (256 elements, 176 bytes):
+//!   [0..2]   d (f16 scale)
+//!   [2..4]   dmin (f16 minimum)
+//!   [4..16]  sc (12-byte packed 6-bit scales+mins for 8 sub-blocks)
+//!   [16..48] qh (32 bytes, 1 high bit per element)
+//!   [48..176] qs (128 bytes of 4-bit low nibbles, 2 per byte)
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+use half::f16;
+
+use super::super::super::dequant_k_quants::unpack_q4k_q5k_scales;
+use super::dot_f32::hsum_f32_neon;
+
+/// Fused dequant+dot for Q5_K using NEON.
+///
+/// # Safety
+/// Requires NEON. Caller must ensure act.len() >= k and blocks covers k/256 blocks.
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+pub unsafe fn fused_dot_q5k_neon(act: &[f32], blocks: &[u8], k: usize) -> f32 {
+    const BLOCK_SIZE: usize = 256;
+    const BLOCK_BYTES: usize = 176;
+    let num_blocks = k / BLOCK_SIZE;
+
+    debug_assert!(act.len() >= k, "act.len() {} < k {}", act.len(), k);
+    debug_assert!(
+        blocks.len() >= num_blocks * BLOCK_BYTES,
+        "blocks.len() {} < required {}",
+        blocks.len(),
+        num_blocks * BLOCK_BYTES
+    );
+
+    let mut total_acc = vdupq_n_f32(0.0);
+
+    for b in 0..num_blocks {
+        let block = &blocks[b * BLOCK_BYTES..];
+        let d = f16::from_le_bytes([block[0], block[1]]).to_f32();
+        let dmin = f16::from_le_bytes([block[2], block[3]]).to_f32();
+        let sc = &block[4..16];
+        let qh = &block[16..48];
+        let qs = &block[48..176];
+        let act_block = &act[b * BLOCK_SIZE..];
+
+        let (scales, mins) = unpack_q4k_q5k_scales(sc);
+
+        for j in 0..8 {
+            let dl = d * scales[j] as f32;
+            let ml = dmin * mins[j] as f32;
+
+            let act_sub = &act_block[j * 32..];
+            let dl_vec = vdupq_n_f32(dl);
+            let ml_vec = vdupq_n_f32(ml);
+
+            // Process 32 elements in 8 groups of 4
+            for g in 0..8 {
+                let l_base = g * 4;
+
+                // Build 5-bit values: low4 from qs + high1 from qh
+                let mut q_vals = [0u32; 4];
+                for l in 0..4 {
+                    let idx = j * 32 + l_base + l;
+                    let qs_idx = j * 16 + (l_base + l) / 2;
+                    let low4 = if (l_base + l) % 2 == 0 {
+                        qs[qs_idx] & 0x0F
+                    } else {
+                        (qs[qs_idx] >> 4) & 0x0F
+                    };
+                    let qh_byte = idx / 8;
+                    let qh_bit = idx % 8;
+                    let high1 = (qh[qh_byte] >> qh_bit) & 0x01;
+                    q_vals[l] = (low4 | (high1 << 4)) as u32;
+                }
+
+                let q_u32 = vcombine_u32(
+                    vcreate_u32(q_vals[0] as u64 | (q_vals[1] as u64) << 32),
+                    vcreate_u32(q_vals[2] as u64 | (q_vals[3] as u64) << 32),
+                );
+                let q_f32 = vcvtq_f32_u32(q_u32);
+
+                let a = vld1q_f32(act_sub.as_ptr().add(l_base));
+
+                // Accumulate: dl * (a * q) - ml * a
+                let aq = vmulq_f32(a, q_f32);
+                total_acc = vfmaq_f32(total_acc, dl_vec, aq);
+                total_acc = vsubq_f32(total_acc, vmulq_f32(ml_vec, a));
+            }
+        }
+    }
+
+    hsum_f32_neon(total_acc)
+}

src/quant/cpu/kernels/simd/aarch64/fused_q6k.rs

@@ -0,0 +1,102 @@
+//! NEON fused dequant+dot for Q6_K format
+//!
+//! Q6_K block layout (256 elements, 210 bytes):
+//!   [0..128]   ql (low 4 bits, packed 2 per byte for 256 elements)
+//!   [128..192]  qh (high 2 bits, packed 4 per byte)
+//!   [192..208]  sc (16 x i8 scales)
+//!   [208..210]  d (f16 scale)
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+use half::f16;
+
+use super::dot_f32::hsum_f32_neon;
+
+/// Fused dequant+dot for Q6_K using NEON.
+///
+/// # Safety
+/// Requires NEON. Caller must ensure act.len() >= k and blocks covers k/256 blocks.
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+pub unsafe fn fused_dot_q6k_neon(act: &[f32], blocks: &[u8], k: usize) -> f32 {
+    const BLOCK_SIZE: usize = 256;
+    const BLOCK_BYTES: usize = 210;
+    let num_blocks = k / BLOCK_SIZE;
+
+    debug_assert!(act.len() >= k, "act.len() {} < k {}", act.len(), k);
+    debug_assert!(
+        blocks.len() >= num_blocks * BLOCK_BYTES,
+        "blocks.len() {} < required {}",
+        blocks.len(),
+        num_blocks * BLOCK_BYTES
+    );
+
+    let mut total_acc = vdupq_n_f32(0.0);
+
+    for b in 0..num_blocks {
+        let block = &blocks[b * BLOCK_BYTES..];
+        let ql = &block[0..128];
+        let qh = &block[128..192];
+        let sc: &[i8] = std::slice::from_raw_parts(block[192..208].as_ptr() as *const i8, 16);
+        let d = f16::from_le_bytes([block[208], block[209]]).to_f32();
+        let act_block = &act[b * BLOCK_SIZE..];
+
+        // Process in two halves of 128 elements
+        for n in 0..2 {
+            let y_base = n * 128;
+            let ql_base = n * 64;
+            let qh_base = n * 32;
+            let sc_base = n * 8;
+
+            // Process 32 elements per sub-iteration, 4 outputs each
+            for l in 0..32 {
+                let is = l / 16;
+
+                // Reconstruct 6-bit values
+                let q1 = ((ql[ql_base + l] & 0x0F) | ((qh[qh_base + l] & 0x03) << 4)) as i8 - 32;
+                let q2 = ((ql[ql_base + l + 32] & 0x0F) | (((qh[qh_base + l] >> 2) & 0x03) << 4))
+                    as i8
+                    - 32;
+                let q3 =
+                    ((ql[ql_base + l] >> 4) | (((qh[qh_base + l] >> 4) & 0x03) << 4)) as i8 - 32;
+                let q4 = ((ql[ql_base + l + 32] >> 4) | (((qh[qh_base + l] >> 6) & 0x03) << 4))
+                    as i8
+                    - 32;
+
+                // Build NEON vector of 4 dequantized values
+                let s1 = d * sc[sc_base + is] as f32;
+                let s2 = d * sc[sc_base + is + 2] as f32;
+                let s3 = d * sc[sc_base + is + 4] as f32;
+                let s4 = d * sc[sc_base + is + 6] as f32;
+
+                let dq = vcombine_f32(
+                    vcreate_f32(
+                        (s1 * q1 as f32).to_bits() as u64
+                            | ((s2 * q2 as f32).to_bits() as u64) << 32,
+                    ),
+                    vcreate_f32(
+                        (s3 * q3 as f32).to_bits() as u64
+                            | ((s4 * q4 as f32).to_bits() as u64) << 32,
+                    ),
+                );
+
+                // Load 4 activation values at positions [y_base+l, y_base+l+32, y_base+l+64, y_base+l+96]
+                let a = vcombine_f32(
+                    vcreate_f32(
+                        (*act_block.as_ptr().add(y_base + l)).to_bits() as u64
+                            | ((*act_block.as_ptr().add(y_base + l + 32)).to_bits() as u64) << 32,
+                    ),
+                    vcreate_f32(
+                        (*act_block.as_ptr().add(y_base + l + 64)).to_bits() as u64
+                            | ((*act_block.as_ptr().add(y_base + l + 96)).to_bits() as u64) << 32,
+                    ),
+                );
+
+                total_acc = vfmaq_f32(total_acc, dq, a);
+            }
+        }
+    }
+
+    hsum_f32_neon(total_acc)
+}

src/quant/cpu/kernels/simd/aarch64/mod.rs

@@ -0,0 +1,4 @@
+pub mod dot_f32;
+pub mod fused_q4k;
+pub mod fused_q5k;
+pub mod fused_q6k;