Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

vendor/ruvector/crates/ruvllm/src/kernels/accelerate.rs

@@ -0,0 +1,793 @@
+//! Apple Accelerate Framework Integration for GEMV
+//!
+//! Provides high-performance matrix-vector multiplication using Apple's
+//! Accelerate framework BLAS implementation. On Apple Silicon, this achieves
+//! significantly higher throughput than hand-written NEON kernels due to:
+//!
+//! - Apple's proprietary AMX (Apple Matrix Extensions) coprocessor
+//! - Highly optimized microarchitecture-specific implementations
+//! - Multi-core parallelization built into the framework
+//!
+//! ## Performance Characteristics (M4 Pro)
+//!
+//! | Operation | NEON Kernel | Accelerate | Speedup |
+//! |-----------|-------------|------------|---------|
+//! | GEMV 4096x4096 | ~35 GFLOPS | ~80+ GFLOPS | ~2.2x |
+//! | GEMV 8192x8192 | ~32 GFLOPS | ~85+ GFLOPS | ~2.7x |
+//!
+//! ## Usage
+//!
+//! The Accelerate backend is automatically selected when:
+//! 1. Running on macOS
+//! 2. The `accelerate` feature is enabled
+//! 3. Matrix dimensions meet minimum thresholds
+//!
+//! ```rust,ignore
+//! use ruvllm::kernels::gemv_accelerate;
+//!
+//! let a = vec![1.0f32; 4096 * 4096];
+//! let x = vec![1.0f32; 4096];
+//! let mut y = vec![0.0f32; 4096];
+//!
+//! // Uses Accelerate framework for optimal performance
+//! gemv_accelerate(&a, &x, &mut y, 4096, 4096, MatrixLayout::RowMajor);
+//! ```
+//!
+//! ## Feature Flag
+//!
+//! Enable with the `accelerate` feature in `Cargo.toml`:
+//! ```toml
+//! ruvllm = { version = "0.1", features = ["accelerate"] }
+//! ```
+
+// ============================================================================
+// FFI Bindings to Apple Accelerate Framework
+// ============================================================================
+
+/// CBLAS matrix storage order
+#[repr(i32)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CblasOrder {
+    /// Row-major storage (C-style)
+    RowMajor = 101,
+    /// Column-major storage (Fortran-style)
+    ColMajor = 102,
+}
+
+/// CBLAS matrix transpose operation
+#[repr(i32)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CblasTranspose {
+    /// No transpose
+    NoTrans = 111,
+    /// Transpose
+    Trans = 112,
+    /// Conjugate transpose (for complex types)
+    ConjTrans = 113,
+}
+
+/// Matrix layout for public API
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum MatrixLayout {
+    /// Row-major storage (C-style) - default for Rust arrays
+    #[default]
+    RowMajor,
+    /// Column-major storage (Fortran-style)
+    ColMajor,
+}
+
+impl From<MatrixLayout> for CblasOrder {
+    fn from(layout: MatrixLayout) -> Self {
+        match layout {
+            MatrixLayout::RowMajor => CblasOrder::RowMajor,
+            MatrixLayout::ColMajor => CblasOrder::ColMajor,
+        }
+    }
+}
+
+// Link against the Accelerate framework on macOS
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+#[link(name = "Accelerate", kind = "framework")]
+extern "C" {
+    /// Single-precision general matrix-vector multiplication
+    ///
+    /// Computes: y = alpha * op(A) * x + beta * y
+    ///
+    /// Where op(A) is either A or A^T depending on `trans`.
+    ///
+    /// # Parameters
+    /// - `order`: Row-major (101) or column-major (102)
+    /// - `trans`: No transpose (111) or transpose (112)
+    /// - `m`: Number of rows of matrix A
+    /// - `n`: Number of columns of matrix A
+    /// - `alpha`: Scalar multiplier for A * x
+    /// - `a`: Pointer to matrix A
+    /// - `lda`: Leading dimension of A (typically n for row-major)
+    /// - `x`: Pointer to vector x
+    /// - `incx`: Increment for x (typically 1)
+    /// - `beta`: Scalar multiplier for y
+    /// - `y`: Pointer to output vector y
+    /// - `incy`: Increment for y (typically 1)
+    fn cblas_sgemv(
+        order: i32,
+        trans: i32,
+        m: i32,
+        n: i32,
+        alpha: f32,
+        a: *const f32,
+        lda: i32,
+        x: *const f32,
+        incx: i32,
+        beta: f32,
+        y: *mut f32,
+        incy: i32,
+    );
+
+    /// Single-precision general matrix-matrix multiplication
+    ///
+    /// Computes: C = alpha * op(A) * op(B) + beta * C
+    fn cblas_sgemm(
+        order: i32,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: f32,
+        a: *const f32,
+        lda: i32,
+        b: *const f32,
+        ldb: i32,
+        beta: f32,
+        c: *mut f32,
+        ldc: i32,
+    );
+
+    /// Single-precision dot product
+    fn cblas_sdot(n: i32, x: *const f32, incx: i32, y: *const f32, incy: i32) -> f32;
+
+    /// Single-precision vector scaling: x = alpha * x
+    fn cblas_sscal(n: i32, alpha: f32, x: *mut f32, incx: i32);
+
+    /// Single-precision axpy: y = alpha * x + y
+    fn cblas_saxpy(n: i32, alpha: f32, x: *const f32, incx: i32, y: *mut f32, incy: i32);
+}
+
+// ============================================================================
+// Public API - Accelerate GEMV
+// ============================================================================
+
+/// Minimum dimension for Accelerate to be beneficial over NEON
+/// Below this threshold, NEON overhead is lower due to function call cost
+const ACCELERATE_MIN_DIM: usize = 256;
+
+/// Minimum total operations (m * n) for Accelerate
+const ACCELERATE_MIN_OPS: usize = 65536; // 256 * 256
+
+/// Check if Accelerate framework is available
+#[inline(always)]
+pub fn is_accelerate_available() -> bool {
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    {
+        true
+    }
+    #[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+    {
+        false
+    }
+}
+
+/// Check if Accelerate should be used for given dimensions
+///
+/// Returns true if:
+/// 1. Accelerate is available
+/// 2. Matrix dimensions are large enough to benefit
+#[inline(always)]
+pub fn should_use_accelerate(m: usize, n: usize) -> bool {
+    is_accelerate_available()
+        && m >= ACCELERATE_MIN_DIM
+        && n >= ACCELERATE_MIN_DIM
+        && m * n >= ACCELERATE_MIN_OPS
+}
+
+/// General Matrix-Vector multiplication using Apple Accelerate
+///
+/// Computes: y = A * x
+///
+/// Uses Apple's BLAS implementation which leverages the AMX coprocessor
+/// on Apple Silicon for maximum throughput.
+///
+/// # Arguments
+/// * `a` - Matrix A (m x n), in specified layout
+/// * `x` - Vector x (n,)
+/// * `y` - Output vector y (m,), modified in-place
+/// * `m` - Number of rows in A
+/// * `n` - Number of columns in A (length of x)
+/// * `layout` - Matrix storage order (RowMajor or ColMajor)
+///
+/// # Performance
+/// On M4 Pro: ~80+ GFLOPS for large matrices (2x+ vs NEON)
+///
+/// # Panics
+/// Panics if dimensions don't match or if not on macOS with accelerate feature
+///
+/// # Example
+/// ```rust,ignore
+/// use ruvllm::kernels::accelerate::{gemv_accelerate, MatrixLayout};
+///
+/// let a = vec![1.0f32; 4096 * 4096];
+/// let x = vec![1.0f32; 4096];
+/// let mut y = vec![0.0f32; 4096];
+///
+/// gemv_accelerate(&a, &x, &mut y, 4096, 4096, MatrixLayout::RowMajor);
+/// ```
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+pub fn gemv_accelerate(
+    a: &[f32],
+    x: &[f32],
+    y: &mut [f32],
+    m: usize,
+    n: usize,
+    layout: MatrixLayout,
+) {
+    debug_assert_eq!(
+        a.len(),
+        m * n,
+        "Matrix A size mismatch: expected {}, got {}",
+        m * n,
+        a.len()
+    );
+    debug_assert_eq!(
+        x.len(),
+        n,
+        "Vector x size mismatch: expected {}, got {}",
+        n,
+        x.len()
+    );
+    debug_assert_eq!(
+        y.len(),
+        m,
+        "Vector y size mismatch: expected {}, got {}",
+        m,
+        y.len()
+    );
+
+    // SECURITY FIX (H-005): Bounds check before i32 cast to prevent overflow
+    // BLAS uses i32 for dimensions, so we must ensure values fit
+    assert!(
+        m <= i32::MAX as usize,
+        "Matrix dimension m={} exceeds i32::MAX for BLAS",
+        m
+    );
+    assert!(
+        n <= i32::MAX as usize,
+        "Matrix dimension n={} exceeds i32::MAX for BLAS",
+        n
+    );
+
+    unsafe {
+        gemv_accelerate_unchecked(a, x, y, m, n, layout);
+    }
+}
+
+/// Unchecked GEMV using Accelerate
+///
+/// # Safety
+/// Caller must ensure:
+/// - `a.len() >= m * n`
+/// - `x.len() >= n`
+/// - `y.len() >= m`
+/// - Pointers are properly aligned
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+#[inline(always)]
+pub unsafe fn gemv_accelerate_unchecked(
+    a: &[f32],
+    x: &[f32],
+    y: &mut [f32],
+    m: usize,
+    n: usize,
+    layout: MatrixLayout,
+) {
+    let order = CblasOrder::from(layout) as i32;
+    let trans = CblasTranspose::NoTrans as i32;
+
+    // For row-major: A is m x n, lda = n
+    // For col-major: A is m x n, lda = m
+    let lda = match layout {
+        MatrixLayout::RowMajor => n as i32,
+        MatrixLayout::ColMajor => m as i32,
+    };
+
+    cblas_sgemv(
+        order,
+        trans,
+        m as i32,
+        n as i32,
+        1.0, // alpha = 1
+        a.as_ptr(),
+        lda,
+        x.as_ptr(),
+        1,   // incx = 1
+        0.0, // beta = 0 (overwrite y)
+        y.as_mut_ptr(),
+        1, // incy = 1
+    );
+}
+
+/// GEMV with transpose using Accelerate
+///
+/// Computes: y = A^T * x
+///
+/// # Arguments
+/// * `a` - Matrix A (m x n), in specified layout
+/// * `x` - Vector x (m,) - note: length is m due to transpose
+/// * `y` - Output vector y (n,), modified in-place
+/// * `m` - Number of rows in A
+/// * `n` - Number of columns in A
+/// * `layout` - Matrix storage order
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+pub fn gemv_transpose_accelerate(
+    a: &[f32],
+    x: &[f32],
+    y: &mut [f32],
+    m: usize,
+    n: usize,
+    layout: MatrixLayout,
+) {
+    debug_assert_eq!(a.len(), m * n);
+    debug_assert_eq!(x.len(), m); // Note: x length is m for transpose
+    debug_assert_eq!(y.len(), n); // Note: y length is n for transpose
+
+    // SECURITY FIX (H-005): Bounds check before i32 cast to prevent overflow
+    assert!(
+        m <= i32::MAX as usize,
+        "Matrix dimension m={} exceeds i32::MAX for BLAS",
+        m
+    );
+    assert!(
+        n <= i32::MAX as usize,
+        "Matrix dimension n={} exceeds i32::MAX for BLAS",
+        n
+    );
+
+    unsafe {
+        let order = CblasOrder::from(layout) as i32;
+        let trans = CblasTranspose::Trans as i32;
+
+        let lda = match layout {
+            MatrixLayout::RowMajor => n as i32,
+            MatrixLayout::ColMajor => m as i32,
+        };
+
+        cblas_sgemv(
+            order,
+            trans,
+            m as i32,
+            n as i32,
+            1.0,
+            a.as_ptr(),
+            lda,
+            x.as_ptr(),
+            1,
+            0.0,
+            y.as_mut_ptr(),
+            1,
+        );
+    }
+}
+
+/// GEMV with alpha and beta scaling using Accelerate
+///
+/// Computes: y = alpha * A * x + beta * y
+///
+/// This is the full BLAS sgemv operation with scaling factors.
+///
+/// # Arguments
+/// * `a` - Matrix A (m x n)
+/// * `x` - Vector x (n,)
+/// * `y` - Vector y (m,), updated in-place
+/// * `m` - Number of rows in A
+/// * `n` - Number of columns in A
+/// * `alpha` - Scalar multiplier for A * x
+/// * `beta` - Scalar multiplier for existing y values
+/// * `layout` - Matrix storage order
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+pub fn gemv_scaled_accelerate(
+    a: &[f32],
+    x: &[f32],
+    y: &mut [f32],
+    m: usize,
+    n: usize,
+    alpha: f32,
+    beta: f32,
+    layout: MatrixLayout,
+) {
+    debug_assert_eq!(a.len(), m * n);
+    debug_assert_eq!(x.len(), n);
+    debug_assert_eq!(y.len(), m);
+
+    // SECURITY FIX (H-005): Bounds check before i32 cast to prevent overflow
+    assert!(
+        m <= i32::MAX as usize,
+        "Matrix dimension m={} exceeds i32::MAX for BLAS",
+        m
+    );
+    assert!(
+        n <= i32::MAX as usize,
+        "Matrix dimension n={} exceeds i32::MAX for BLAS",
+        n
+    );
+
+    unsafe {
+        let order = CblasOrder::from(layout) as i32;
+        let trans = CblasTranspose::NoTrans as i32;
+
+        let lda = match layout {
+            MatrixLayout::RowMajor => n as i32,
+            MatrixLayout::ColMajor => m as i32,
+        };
+
+        cblas_sgemv(
+            order,
+            trans,
+            m as i32,
+            n as i32,
+            alpha,
+            a.as_ptr(),
+            lda,
+            x.as_ptr(),
+            1,
+            beta,
+            y.as_mut_ptr(),
+            1,
+        );
+    }
+}
+
+// ============================================================================
+// Public API - Accelerate GEMM
+// ============================================================================
+
+/// General Matrix-Matrix multiplication using Apple Accelerate
+///
+/// Computes: C = A * B
+///
+/// # Arguments
+/// * `a` - Matrix A (m x k), row-major
+/// * `b` - Matrix B (k x n), row-major
+/// * `c` - Output matrix C (m x n), row-major, modified in-place
+/// * `m` - Number of rows in A and C
+/// * `k` - Number of columns in A, rows in B
+/// * `n` - Number of columns in B and C
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+pub fn gemm_accelerate(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
+    debug_assert_eq!(a.len(), m * k);
+    debug_assert_eq!(b.len(), k * n);
+    debug_assert_eq!(c.len(), m * n);
+
+    // SECURITY FIX (H-005): Bounds check before i32 cast to prevent overflow
+    assert!(
+        m <= i32::MAX as usize,
+        "Matrix dimension m={} exceeds i32::MAX for BLAS",
+        m
+    );
+    assert!(
+        k <= i32::MAX as usize,
+        "Matrix dimension k={} exceeds i32::MAX for BLAS",
+        k
+    );
+    assert!(
+        n <= i32::MAX as usize,
+        "Matrix dimension n={} exceeds i32::MAX for BLAS",
+        n
+    );
+
+    unsafe {
+        cblas_sgemm(
+            CblasOrder::RowMajor as i32,
+            CblasTranspose::NoTrans as i32,
+            CblasTranspose::NoTrans as i32,
+            m as i32,
+            n as i32,
+            k as i32,
+            1.0, // alpha
+            a.as_ptr(),
+            k as i32, // lda
+            b.as_ptr(),
+            n as i32, // ldb
+            0.0,      // beta
+            c.as_mut_ptr(),
+            n as i32, // ldc
+        );
+    }
+}
+
+// ============================================================================
+// Additional BLAS Operations
+// ============================================================================
+
+/// Single-precision dot product using Accelerate
+///
+/// Computes: result = x . y
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+#[inline]
+pub fn dot_accelerate(x: &[f32], y: &[f32]) -> f32 {
+    debug_assert_eq!(x.len(), y.len());
+    unsafe { cblas_sdot(x.len() as i32, x.as_ptr(), 1, y.as_ptr(), 1) }
+}
+
+/// Scale vector in-place using Accelerate
+///
+/// Computes: x = alpha * x
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+#[inline]
+pub fn scal_accelerate(x: &mut [f32], alpha: f32) {
+    unsafe { cblas_sscal(x.len() as i32, alpha, x.as_mut_ptr(), 1) }
+}
+
+/// Vector addition with scaling using Accelerate
+///
+/// Computes: y = alpha * x + y
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+#[inline]
+pub fn axpy_accelerate(x: &[f32], y: &mut [f32], alpha: f32) {
+    debug_assert_eq!(x.len(), y.len());
+    unsafe { cblas_saxpy(x.len() as i32, alpha, x.as_ptr(), 1, y.as_mut_ptr(), 1) }
+}
+
+// ============================================================================
+// Fallback implementations for non-macOS platforms
+// ============================================================================
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn gemv_accelerate(
+    _a: &[f32],
+    _x: &[f32],
+    _y: &mut [f32],
+    _m: usize,
+    _n: usize,
+    _layout: MatrixLayout,
+) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub unsafe fn gemv_accelerate_unchecked(
+    _a: &[f32],
+    _x: &[f32],
+    _y: &mut [f32],
+    _m: usize,
+    _n: usize,
+    _layout: MatrixLayout,
+) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn gemv_transpose_accelerate(
+    _a: &[f32],
+    _x: &[f32],
+    _y: &mut [f32],
+    _m: usize,
+    _n: usize,
+    _layout: MatrixLayout,
+) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn gemv_scaled_accelerate(
+    _a: &[f32],
+    _x: &[f32],
+    _y: &mut [f32],
+    _m: usize,
+    _n: usize,
+    _alpha: f32,
+    _beta: f32,
+    _layout: MatrixLayout,
+) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn gemm_accelerate(_a: &[f32], _b: &[f32], _c: &mut [f32], _m: usize, _k: usize, _n: usize) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn dot_accelerate(_x: &[f32], _y: &[f32]) -> f32 {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn scal_accelerate(_x: &mut [f32], _alpha: f32) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn axpy_accelerate(_x: &[f32], _y: &mut [f32], _alpha: f32) {
+    panic!("Accelerate framework is only available on macOS with 'accelerate' feature enabled");
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_accelerate_availability() {
+        #[cfg(all(target_os = "macos", feature = "accelerate"))]
+        assert!(is_accelerate_available());
+
+        #[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+        assert!(!is_accelerate_available());
+    }
+
+    #[test]
+    fn test_should_use_accelerate_thresholds() {
+        // Below threshold
+        assert!(!should_use_accelerate(128, 128));
+        assert!(!should_use_accelerate(255, 256));
+
+        // At/above threshold (only true on macOS with feature)
+        #[cfg(all(target_os = "macos", feature = "accelerate"))]
+        {
+            assert!(should_use_accelerate(256, 256));
+            assert!(should_use_accelerate(4096, 4096));
+        }
+
+        #[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+        {
+            assert!(!should_use_accelerate(256, 256));
+            assert!(!should_use_accelerate(4096, 4096));
+        }
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_gemv_accelerate_correctness() {
+        // Simple 2x3 matrix test
+        // A = [[1, 2, 3],
+        //      [4, 5, 6]]
+        // x = [1, 1, 1]
+        // y = A * x = [6, 15]
+        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let x = vec![1.0, 1.0, 1.0];
+        let mut y = vec![0.0, 0.0];
+
+        gemv_accelerate(&a, &x, &mut y, 2, 3, MatrixLayout::RowMajor);
+
+        assert!((y[0] - 6.0).abs() < 1e-5);
+        assert!((y[1] - 15.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_gemv_transpose_correctness() {
+        // A = [[1, 2, 3],
+        //      [4, 5, 6]]
+        // x = [1, 1]
+        // y = A^T * x = [5, 7, 9]
+        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let x = vec![1.0, 1.0];
+        let mut y = vec![0.0, 0.0, 0.0];
+
+        gemv_transpose_accelerate(&a, &x, &mut y, 2, 3, MatrixLayout::RowMajor);
+
+        assert!((y[0] - 5.0).abs() < 1e-5);
+        assert!((y[1] - 7.0).abs() < 1e-5);
+        assert!((y[2] - 9.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_gemv_scaled_correctness() {
+        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let x = vec![1.0, 1.0, 1.0];
+        let mut y = vec![1.0, 2.0]; // Initial values
+
+        // y = 2 * A * x + 3 * y
+        // y = 2 * [6, 15] + 3 * [1, 2] = [12, 30] + [3, 6] = [15, 36]
+        gemv_scaled_accelerate(&a, &x, &mut y, 2, 3, 2.0, 3.0, MatrixLayout::RowMajor);
+
+        assert!((y[0] - 15.0).abs() < 1e-5);
+        assert!((y[1] - 36.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_gemm_accelerate_correctness() {
+        // A = [[1, 2],
+        //      [3, 4]]
+        // B = [[5, 6],
+        //      [7, 8]]
+        // C = A * B = [[19, 22],
+        //              [43, 50]]
+        let a = vec![1.0, 2.0, 3.0, 4.0];
+        let b = vec![5.0, 6.0, 7.0, 8.0];
+        let mut c = vec![0.0; 4];
+
+        gemm_accelerate(&a, &b, &mut c, 2, 2, 2);
+
+        assert!((c[0] - 19.0).abs() < 1e-5);
+        assert!((c[1] - 22.0).abs() < 1e-5);
+        assert!((c[2] - 43.0).abs() < 1e-5);
+        assert!((c[3] - 50.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_dot_accelerate_correctness() {
+        let x = vec![1.0, 2.0, 3.0];
+        let y = vec![4.0, 5.0, 6.0];
+
+        let result = dot_accelerate(&x, &y);
+
+        // 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32
+        assert!((result - 32.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_scal_accelerate_correctness() {
+        let mut x = vec![1.0, 2.0, 3.0];
+
+        scal_accelerate(&mut x, 2.0);
+
+        assert!((x[0] - 2.0).abs() < 1e-5);
+        assert!((x[1] - 4.0).abs() < 1e-5);
+        assert!((x[2] - 6.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_axpy_accelerate_correctness() {
+        let x = vec![1.0, 2.0, 3.0];
+        let mut y = vec![4.0, 5.0, 6.0];
+
+        // y = 2 * x + y = [2, 4, 6] + [4, 5, 6] = [6, 9, 12]
+        axpy_accelerate(&x, &mut y, 2.0);
+
+        assert!((y[0] - 6.0).abs() < 1e-5);
+        assert!((y[1] - 9.0).abs() < 1e-5);
+        assert!((y[2] - 12.0).abs() < 1e-5);
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_gemv_large_matrix() {
+        // Test with a larger matrix to verify performance path
+        let m = 512;
+        let n = 512;
+        let a: Vec<f32> = (0..m * n).map(|i| (i % 10) as f32 * 0.1).collect();
+        let x: Vec<f32> = vec![1.0; n];
+        let mut y = vec![0.0; m];
+
+        gemv_accelerate(&a, &x, &mut y, m, n, MatrixLayout::RowMajor);
+
+        // Verify non-zero results
+        assert!(y.iter().any(|&v| v != 0.0));
+    }
+
+    #[cfg(all(target_os = "macos", feature = "accelerate"))]
+    #[test]
+    fn test_col_major_layout() {
+        // Test column-major layout
+        // A stored col-major: column 0 = [1, 4], column 1 = [2, 5], column 2 = [3, 6]
+        // Storage: [1, 4, 2, 5, 3, 6]
+        // Logical matrix (2x3):
+        // [[1, 2, 3],
+        //  [4, 5, 6]]
+        let a = vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]; // Column-major storage
+        let x = vec![1.0, 1.0, 1.0];
+        let mut y = vec![0.0, 0.0];
+
+        gemv_accelerate(&a, &x, &mut y, 2, 3, MatrixLayout::ColMajor);
+
+        assert!((y[0] - 6.0).abs() < 1e-5);
+        assert!((y[1] - 15.0).abs() < 1e-5);
+    }
+}

vendor/ruvector/crates/ruvllm/src/kernels/mod.rs

@@ -0,0 +1,312 @@
+//! NEON-Optimized LLM Kernels for Mac M4 Pro
+//!
+//! This module provides highly optimized SIMD kernels for LLM operations,
+//! specifically tuned for Apple Silicon (M1/M2/M3/M4) using ARM NEON intrinsics.
+//!
+//! ## Quick Start
+//!
+//! ```rust,ignore
+//! use ruvllm::kernels::{
+//!     flash_attention_neon, apply_rope_neon, rms_norm_neon,
+//!     AttentionConfig, is_neon_available,
+//! };
+//!
+//! // Check NEON availability
+//! assert!(is_neon_available(), "NEON required for optimal performance");
+//!
+//! // Configure attention
+//! let config = AttentionConfig {
+//!     num_heads: 32,
+//!     num_kv_heads: 8,  // GQA with 4:1 ratio
+//!     head_dim: 128,
+//!     causal: true,
+//!     ..Default::default()
+//! };
+//!
+//! // Flash attention with NEON SIMD
+//! let output = flash_attention_neon(
+//!     &query, &key, &value,
+//!     config.effective_scale(),
+//!     config.causal
+//! );
+//!
+//! // Apply RoPE to query/key tensors
+//! apply_rope_neon(&mut qk, &positions, config.head_dim, 10000.0);
+//!
+//! // RMSNorm normalization
+//! rms_norm_neon(&mut hidden, &weight, 1e-6);
+//! ```
+//!
+//! ## Kernel Categories
+//!
+//! - [`attention`]: Flash Attention 2, Paged Attention, MQA/GQA
+//! - [`rope`]: Rotary Position Embeddings (RoPE)
+//! - [`norm`]: RMSNorm, LayerNorm
+//! - [`matmul`]: Batched GEMM operations
+//! - [`quantized`]: INT8/INT4 quantized inference kernels
+//! - [`activations`]: Vectorized SiLU, GELU, ReLU, Softmax
+//!
+//! ## Performance Characteristics
+//!
+//! | Kernel | Sequence Length | Throughput | vs. Naive |
+//! |--------|-----------------|------------|-----------|
+//! | `flash_attention_neon` | 4096 | 2.5 GFLOPS | 3.2x |
+//! | `paged_attention_neon` | 8192+ | 2.1 GFLOPS | 2.8x |
+//! | `rms_norm_neon` | Any | 4.8 GFLOPS | 4.1x |
+//! | `gemm_neon` | 4096x4096 | 1.2 GFLOPS | 2.4x |
+//! | `silu` | Any | 5.2 GFLOPS | 3.5x |
+//! | `gelu` | Any | 4.5 GFLOPS | 3.2x |
+//! | `softmax` | Any | 3.8 GFLOPS | 2.8x |
+//!
+//! ## Performance Optimizations
+//!
+//! All kernels implement:
+//! - 4x loop unrolling for instruction-level parallelism
+//! - FMA instructions for improved throughput
+//! - Pointer caching to reduce address calculations
+//! - Efficient horizontal reductions via `vaddvq_f32`
+//! - Software prefetching for large tensors
+//!
+//! ## Memory Layout
+//!
+//! Kernels expect contiguous memory in the following layouts:
+//!
+//! - **Query/Key/Value**: `[batch, seq_len, num_heads, head_dim]`
+//! - **KV Cache**: `[batch, num_kv_heads, seq_len, head_dim]`
+//! - **Hidden states**: `[batch, seq_len, hidden_dim]`
+
+pub mod activations;
+pub mod attention;
+pub mod matmul;
+pub mod norm;
+pub mod quantized;
+pub mod rope;
+
+// Apple Accelerate framework integration (macOS only)
+#[cfg(any(target_os = "macos", doc))]
+pub mod accelerate;
+
+// Apple Neural Engine (ANE) optimized operations
+// Uses BNNS (Basic Neural Network Subroutines) which routes to ANE on macOS.
+// Decision-logic functions (should_use_ane, get_ane_recommendation, etc.) and
+// platform-fallback stubs are available on all targets so that tests and
+// cross-platform code can reference the module unconditionally.
+pub mod ane_ops;
+
+// Re-exports for convenience
+pub use attention::{
+    flash_attention_auto,
+    // TD-009: Zero-allocation attention functions and scratch buffers
+    flash_attention_into,
+    flash_attention_neon,
+    flash_attention_v2,
+    flash_attention_with_scratch,
+    grouped_query_attention_neon,
+    multi_query_attention_neon,
+    paged_attention_neon,
+    select_block_size,
+    AttentionScratch,
+    PagedKvCache,
+    BLOCK_SIZE_LARGE,
+    BLOCK_SIZE_MEDIUM,
+    BLOCK_SIZE_SMALL,
+};
+// Thread-local scratch buffer for zero-allocation attention (non-WASM only)
+#[cfg(not(target_arch = "wasm32"))]
+pub use attention::THREAD_LOCAL_SCRATCH;
+#[cfg(all(feature = "parallel", not(target_arch = "wasm32")))]
+pub use attention::{
+    grouped_query_attention_parallel, multi_head_attention_parallel, multi_query_attention_parallel,
+};
+pub use matmul::{batched_gemm_neon, gemm_neon, gemv_neon};
+#[cfg(all(feature = "parallel", not(target_arch = "wasm32")))]
+pub use matmul::{
+    batched_gemm_parallel, configure_thread_pool, gemm_parallel, gemv_parallel, get_physical_cores,
+};
+pub use norm::{layer_norm_neon, rms_norm_neon};
+pub use quantized::{
+    dequantize_int4, dequantize_int8, int4_gemv_neon, int8_gemv_neon, q4k_gemv_neon,
+    quantize_to_int4, quantize_to_int8, quantize_to_q4k, BlockQ4K, QuantizedInt4, QuantizedInt8,
+    INT4_BLOCK_SIZE, Q4K_SUPER_BLOCK_SIZE,
+};
+pub use rope::{apply_rope_neon, precompute_rope_tables, RopeConfig};
+
+// Activation function exports
+pub use activations::{
+    batch_gelu, batch_silu, batch_softmax, gelu, gelu_exact, gelu_vec, leaky_relu, relu, relu_vec,
+    silu, silu_vec, softmax, softmax_temperature, softmax_vec,
+};
+
+// Accelerate framework exports (macOS only)
+#[cfg(all(target_os = "macos", feature = "accelerate"))]
+pub use accelerate::{
+    axpy_accelerate, dot_accelerate, gemm_accelerate, gemv_accelerate, gemv_scaled_accelerate,
+    gemv_transpose_accelerate, is_accelerate_available, scal_accelerate, should_use_accelerate,
+    MatrixLayout,
+};
+
+// Fallback availability check for non-macOS platforms
+#[cfg(not(all(target_os = "macos", feature = "accelerate")))]
+pub fn is_accelerate_available() -> bool {
+    false
+}
+
+// ANE (Apple Neural Engine) ops exports (macOS only with coreml feature)
+#[cfg(all(target_os = "macos", feature = "coreml"))]
+pub use ane_ops::{
+    batched_matmul_ane,
+    gelu_ane,
+    gelu_auto,
+    // Strategy recommendations (M4 Pro optimized)
+    get_ane_recommendation,
+    // Availability checks
+    is_ane_available,
+    layer_norm_ane,
+    layer_norm_auto,
+    // Direct ANE operations
+    matmul_ane,
+    // Auto-dispatch functions
+    matmul_auto,
+    rms_norm_ane,
+    rms_norm_auto,
+    should_use_ane,
+    should_use_ane_activation,
+    should_use_ane_matmul,
+    silu_ane,
+    silu_auto,
+    softmax_ane,
+    softmax_auto,
+    AneRecommendation,
+};
+
+// Re-export ANE availability check for all platforms without coreml feature
+// (the ane_ops module is now unconditionally available)
+#[cfg(not(all(target_os = "macos", feature = "coreml")))]
+pub use ane_ops::is_ane_available;
+
+/// SIMD lane width for NEON (128-bit = 4 floats).
+///
+/// ARM NEON registers are 128 bits wide, holding 4 single-precision floats.
+/// This constant is used for loop unrolling and vectorization decisions.
+pub const NEON_LANE_WIDTH: usize = 4;
+
+/// Optimal unroll factor for M4 Pro's 6-wide superscalar core.
+///
+/// The M4 Pro can execute up to 6 operations per cycle. Using a 4x unroll
+/// factor with FMA instructions achieves near-optimal utilization.
+pub const UNROLL_FACTOR: usize = 4;
+
+/// Prefetch distance in cache lines (64 bytes = 16 floats)
+pub const PREFETCH_DISTANCE: usize = 64;
+
+/// Check if NEON is available at runtime
+#[inline(always)]
+pub fn is_neon_available() -> bool {
+    #[cfg(target_arch = "aarch64")]
+    {
+        true // NEON is always available on aarch64
+    }
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        false
+    }
+}
+
+/// Kernel configuration for attention operations.
+///
+/// Configures the attention mechanism including head counts, dimensions,
+/// and masking behavior. Supports both standard multi-head attention and
+/// grouped-query attention (GQA).
+///
+/// # Example
+///
+/// ```rust,ignore
+/// use ruvllm::kernels::AttentionConfig;
+///
+/// // Standard Mistral-7B configuration with GQA
+/// let config = AttentionConfig {
+///     num_heads: 32,
+///     num_kv_heads: 8,  // 4:1 GQA ratio
+///     head_dim: 128,
+///     max_seq_len: 4096,
+///     causal: true,
+///     scale: 0.0,  // Auto-computed as 1/sqrt(head_dim)
+/// };
+///
+/// assert_eq!(config.gqa_ratio(), 4);
+/// assert!((config.effective_scale() - 0.0884).abs() < 0.001);
+/// ```
+///
+/// # GQA (Grouped-Query Attention)
+///
+/// GQA reduces memory usage by sharing key-value heads across query heads:
+///
+/// | GQA Ratio | KV Memory | Quality |
+/// |-----------|-----------|---------|
+/// | 1:1 (MHA) | 100% | Best |
+/// | 4:1 | 25% | Excellent |
+/// | 8:1 | 12.5% | Good |
+#[derive(Debug, Clone, Copy)]
+pub struct AttentionConfig {
+    /// Number of query heads
+    pub num_heads: usize,
+    /// Number of key-value heads (for GQA)
+    pub num_kv_heads: usize,
+    /// Dimension per head
+    pub head_dim: usize,
+    /// Maximum sequence length
+    pub max_seq_len: usize,
+    /// Whether to use causal masking
+    pub causal: bool,
+    /// Softmax scale factor (typically 1/sqrt(head_dim))
+    pub scale: f32,
+}
+
+impl Default for AttentionConfig {
+    fn default() -> Self {
+        Self {
+            num_heads: 32,
+            num_kv_heads: 8,
+            head_dim: 128,
+            max_seq_len: 4096,
+            causal: true,
+            scale: 0.0, // Will be computed from head_dim if 0
+        }
+    }
+}
+
+impl AttentionConfig {
+    /// Get the effective scale (computes from head_dim if not set)
+    #[inline(always)]
+    pub fn effective_scale(&self) -> f32 {
+        if self.scale == 0.0 {
+            1.0 / (self.head_dim as f32).sqrt()
+        } else {
+            self.scale
+        }
+    }
+
+    /// Get the GQA ratio (num_heads / num_kv_heads)
+    #[inline(always)]
+    pub fn gqa_ratio(&self) -> usize {
+        self.num_heads / self.num_kv_heads
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_attention_config() {
+        let config = AttentionConfig::default();
+        assert_eq!(config.gqa_ratio(), 4);
+        assert!((config.effective_scale() - 0.088388).abs() < 0.001);
+    }
+
+    #[test]
+    fn test_neon_available() {
+        #[cfg(target_arch = "aarch64")]
+        assert!(is_neon_available());
+    }
+}

vendor/ruvector/crates/ruvllm/src/kernels/norm.rs

@@ -0,0 +1,651 @@
+//! NEON-Optimized Normalization Layers
+//!
+//! Implements efficient normalization operations for transformer models:
+//!
+//! - **RMSNorm**: Root Mean Square normalization (Llama, Mistral)
+//! - **LayerNorm**: Standard layer normalization (GPT, BERT)
+//! - **GroupNorm**: Group normalization (Vision models)
+//!
+//! ## Performance Characteristics
+//!
+//! | Operation | Dimension | M4 Pro Throughput |
+//! |-----------|-----------|-------------------|
+//! | RMSNorm | 4096 | ~12 GB/s |
+//! | LayerNorm | 4096 | ~10 GB/s |
+//! | GroupNorm | 4096 | ~8 GB/s |
+//!
+//! ## Why RMSNorm?
+//!
+//! RMSNorm is faster than LayerNorm because:
+//! 1. No mean computation (saves one reduction)
+//! 2. No mean subtraction (saves one element-wise op)
+//! 3. Simpler gradient computation
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+use super::{NEON_LANE_WIDTH, UNROLL_FACTOR};
+
+/// RMSNorm with NEON optimization
+///
+/// Applies Root Mean Square normalization:
+/// ```text
+/// output = x * weight / sqrt(mean(x^2) + eps)
+/// ```
+///
+/// # Arguments
+/// * `x` - Input tensor (modified in-place)
+/// * `weight` - Learnable scale parameters
+/// * `eps` - Small constant for numerical stability
+///
+/// # Panics
+/// Panics if `x.len() != weight.len()`
+#[inline(always)]
+pub fn rms_norm_neon(x: &mut [f32], weight: &[f32], eps: f32) {
+    debug_assert_eq!(x.len(), weight.len());
+
+    let len = x.len();
+    if len == 0 {
+        return;
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        rms_norm_neon_impl(x, weight, eps);
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        rms_norm_scalar(x, weight, eps);
+    }
+}
+
+/// NEON implementation of RMSNorm
+#[cfg(target_arch = "aarch64")]
+#[inline(always)]
+unsafe fn rms_norm_neon_impl(x: &mut [f32], weight: &[f32], eps: f32) {
+    let len = x.len();
+    let x_ptr = x.as_mut_ptr();
+    let w_ptr = weight.as_ptr();
+
+    // Step 1: Compute sum of squares using 4x unrolling
+    let mut sum0 = vdupq_n_f32(0.0);
+    let mut sum1 = vdupq_n_f32(0.0);
+    let mut sum2 = vdupq_n_f32(0.0);
+    let mut sum3 = vdupq_n_f32(0.0);
+
+    let chunks = len / (NEON_LANE_WIDTH * UNROLL_FACTOR);
+    let mut idx = 0usize;
+
+    for _ in 0..chunks {
+        let v0 = vld1q_f32(x_ptr.add(idx));
+        sum0 = vfmaq_f32(sum0, v0, v0);
+
+        let v1 = vld1q_f32(x_ptr.add(idx + 4));
+        sum1 = vfmaq_f32(sum1, v1, v1);
+
+        let v2 = vld1q_f32(x_ptr.add(idx + 8));
+        sum2 = vfmaq_f32(sum2, v2, v2);
+
+        let v3 = vld1q_f32(x_ptr.add(idx + 12));
+        sum3 = vfmaq_f32(sum3, v3, v3);
+
+        idx += 16;
+    }
+
+    // Combine accumulators
+    let sum01 = vaddq_f32(sum0, sum1);
+    let sum23 = vaddq_f32(sum2, sum3);
+    let sum = vaddq_f32(sum01, sum23);
+
+    // Process remaining 4-element chunks
+    let remaining_chunks = (len - idx) / NEON_LANE_WIDTH;
+    let mut final_sum = sum;
+    for _ in 0..remaining_chunks {
+        let v = vld1q_f32(x_ptr.add(idx));
+        final_sum = vfmaq_f32(final_sum, v, v);
+        idx += 4;
+    }
+
+    // Horizontal sum
+    let mut sum_sq = vaddvq_f32(final_sum);
+
+    // Handle remaining elements
+    for i in idx..len {
+        let v = *x_ptr.add(i);
+        sum_sq += v * v;
+    }
+
+    // Step 2: Compute normalization factor
+    let mean_sq = sum_sq / len as f32;
+    let rms = (mean_sq + eps).sqrt();
+    let inv_rms = 1.0 / rms;
+    let inv_rms_vec = vdupq_n_f32(inv_rms);
+
+    // Step 3: Apply normalization and weight with 4x unrolling
+    idx = 0;
+    for _ in 0..chunks {
+        let x0 = vld1q_f32(x_ptr.add(idx));
+        let w0 = vld1q_f32(w_ptr.add(idx));
+        vst1q_f32(x_ptr.add(idx), vmulq_f32(vmulq_f32(x0, inv_rms_vec), w0));
+
+        let x1 = vld1q_f32(x_ptr.add(idx + 4));
+        let w1 = vld1q_f32(w_ptr.add(idx + 4));
+        vst1q_f32(
+            x_ptr.add(idx + 4),
+            vmulq_f32(vmulq_f32(x1, inv_rms_vec), w1),
+        );
+
+        let x2 = vld1q_f32(x_ptr.add(idx + 8));
+        let w2 = vld1q_f32(w_ptr.add(idx + 8));
+        vst1q_f32(
+            x_ptr.add(idx + 8),
+            vmulq_f32(vmulq_f32(x2, inv_rms_vec), w2),
+        );
+
+        let x3 = vld1q_f32(x_ptr.add(idx + 12));
+        let w3 = vld1q_f32(w_ptr.add(idx + 12));
+        vst1q_f32(
+            x_ptr.add(idx + 12),
+            vmulq_f32(vmulq_f32(x3, inv_rms_vec), w3),
+        );
+
+        idx += 16;
+    }
+
+    // Remaining chunks
+    for _ in 0..remaining_chunks {
+        let x_v = vld1q_f32(x_ptr.add(idx));
+        let w_v = vld1q_f32(w_ptr.add(idx));
+        vst1q_f32(x_ptr.add(idx), vmulq_f32(vmulq_f32(x_v, inv_rms_vec), w_v));
+        idx += 4;
+    }
+
+    // Remaining elements
+    for i in idx..len {
+        *x_ptr.add(i) = *x_ptr.add(i) * inv_rms * *w_ptr.add(i);
+    }
+}
+
+/// Scalar fallback for RMSNorm
+#[allow(dead_code)]
+fn rms_norm_scalar(x: &mut [f32], weight: &[f32], eps: f32) {
+    let len = x.len();
+
+    // Compute sum of squares
+    let sum_sq: f32 = x.iter().map(|v| v * v).sum();
+
+    // Compute normalization factor
+    let mean_sq = sum_sq / len as f32;
+    let inv_rms = 1.0 / (mean_sq + eps).sqrt();
+
+    // Apply normalization and weight
+    for (i, w) in weight.iter().enumerate() {
+        x[i] = x[i] * inv_rms * w;
+    }
+}
+
+/// LayerNorm with NEON optimization
+///
+/// Applies Layer normalization:
+/// ```text
+/// output = (x - mean) / sqrt(var + eps) * weight + bias
+/// ```
+///
+/// # Arguments
+/// * `x` - Input tensor (modified in-place)
+/// * `weight` - Learnable scale parameters (gamma)
+/// * `bias` - Learnable shift parameters (beta)
+/// * `eps` - Small constant for numerical stability
+///
+/// # Panics
+/// Panics if `x.len() != weight.len() || x.len() != bias.len()`
+#[inline(always)]
+pub fn layer_norm_neon(x: &mut [f32], weight: &[f32], bias: &[f32], eps: f32) {
+    debug_assert_eq!(x.len(), weight.len());
+    debug_assert_eq!(x.len(), bias.len());
+
+    let len = x.len();
+    if len == 0 {
+        return;
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        layer_norm_neon_impl(x, weight, bias, eps);
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        layer_norm_scalar(x, weight, bias, eps);
+    }
+}
+
+/// NEON implementation of LayerNorm
+#[cfg(target_arch = "aarch64")]
+#[inline(always)]
+unsafe fn layer_norm_neon_impl(x: &mut [f32], weight: &[f32], bias: &[f32], eps: f32) {
+    let len = x.len();
+    let x_ptr = x.as_mut_ptr();
+    let w_ptr = weight.as_ptr();
+    let b_ptr = bias.as_ptr();
+
+    // Step 1: Compute sum (for mean) and sum of squares using 4x unrolling
+    let mut sum0 = vdupq_n_f32(0.0);
+    let mut sum1 = vdupq_n_f32(0.0);
+    let mut sq0 = vdupq_n_f32(0.0);
+    let mut sq1 = vdupq_n_f32(0.0);
+
+    let chunks = len / (NEON_LANE_WIDTH * 2);
+    let mut idx = 0usize;
+
+    for _ in 0..chunks {
+        let v0 = vld1q_f32(x_ptr.add(idx));
+        sum0 = vaddq_f32(sum0, v0);
+        sq0 = vfmaq_f32(sq0, v0, v0);
+
+        let v1 = vld1q_f32(x_ptr.add(idx + 4));
+        sum1 = vaddq_f32(sum1, v1);
+        sq1 = vfmaq_f32(sq1, v1, v1);
+
+        idx += 8;
+    }
+
+    // Combine
+    let sum_vec = vaddq_f32(sum0, sum1);
+    let sq_vec = vaddq_f32(sq0, sq1);
+
+    // Process remaining chunks
+    let remaining_chunks = (len - idx) / NEON_LANE_WIDTH;
+    let mut final_sum = sum_vec;
+    let mut final_sq = sq_vec;
+    for _ in 0..remaining_chunks {
+        let v = vld1q_f32(x_ptr.add(idx));
+        final_sum = vaddq_f32(final_sum, v);
+        final_sq = vfmaq_f32(final_sq, v, v);
+        idx += 4;
+    }
+
+    // Horizontal sums
+    let mut sum = vaddvq_f32(final_sum);
+    let mut sum_sq = vaddvq_f32(final_sq);
+
+    // Handle remaining elements
+    for i in idx..len {
+        let v = *x_ptr.add(i);
+        sum += v;
+        sum_sq += v * v;
+    }
+
+    // Step 2: Compute mean and variance
+    let n = len as f32;
+    let mean = sum / n;
+    let variance = (sum_sq / n) - (mean * mean);
+    let inv_std = 1.0 / (variance + eps).sqrt();
+
+    let mean_vec = vdupq_n_f32(mean);
+    let inv_std_vec = vdupq_n_f32(inv_std);
+
+    // Step 3: Apply normalization, weight, and bias with 4x unrolling
+    idx = 0;
+    let unroll_chunks = len / (NEON_LANE_WIDTH * UNROLL_FACTOR);
+    for _ in 0..unroll_chunks {
+        // Normalize: (x - mean) * inv_std
+        let x0 = vld1q_f32(x_ptr.add(idx));
+        let n0 = vmulq_f32(vsubq_f32(x0, mean_vec), inv_std_vec);
+        let w0 = vld1q_f32(w_ptr.add(idx));
+        let b0 = vld1q_f32(b_ptr.add(idx));
+        vst1q_f32(x_ptr.add(idx), vfmaq_f32(b0, n0, w0));
+
+        let x1 = vld1q_f32(x_ptr.add(idx + 4));
+        let n1 = vmulq_f32(vsubq_f32(x1, mean_vec), inv_std_vec);
+        let w1 = vld1q_f32(w_ptr.add(idx + 4));
+        let b1 = vld1q_f32(b_ptr.add(idx + 4));
+        vst1q_f32(x_ptr.add(idx + 4), vfmaq_f32(b1, n1, w1));
+
+        let x2 = vld1q_f32(x_ptr.add(idx + 8));
+        let n2 = vmulq_f32(vsubq_f32(x2, mean_vec), inv_std_vec);
+        let w2 = vld1q_f32(w_ptr.add(idx + 8));
+        let b2 = vld1q_f32(b_ptr.add(idx + 8));
+        vst1q_f32(x_ptr.add(idx + 8), vfmaq_f32(b2, n2, w2));
+
+        let x3 = vld1q_f32(x_ptr.add(idx + 12));
+        let n3 = vmulq_f32(vsubq_f32(x3, mean_vec), inv_std_vec);
+        let w3 = vld1q_f32(w_ptr.add(idx + 12));
+        let b3 = vld1q_f32(b_ptr.add(idx + 12));
+        vst1q_f32(x_ptr.add(idx + 12), vfmaq_f32(b3, n3, w3));
+
+        idx += 16;
+    }
+
+    // Remaining chunks
+    let remaining = (len - idx) / NEON_LANE_WIDTH;
+    for _ in 0..remaining {
+        let x_v = vld1q_f32(x_ptr.add(idx));
+        let n_v = vmulq_f32(vsubq_f32(x_v, mean_vec), inv_std_vec);
+        let w_v = vld1q_f32(w_ptr.add(idx));
+        let b_v = vld1q_f32(b_ptr.add(idx));
+        vst1q_f32(x_ptr.add(idx), vfmaq_f32(b_v, n_v, w_v));
+        idx += 4;
+    }
+
+    // Remaining elements
+    for i in idx..len {
+        let normalized = (*x_ptr.add(i) - mean) * inv_std;
+        *x_ptr.add(i) = normalized * *w_ptr.add(i) + *b_ptr.add(i);
+    }
+}
+
+/// Scalar fallback for LayerNorm
+#[allow(dead_code)]
+fn layer_norm_scalar(x: &mut [f32], weight: &[f32], bias: &[f32], eps: f32) {
+    let len = x.len();
+    let n = len as f32;
+
+    // Compute mean
+    let sum: f32 = x.iter().sum();
+    let mean = sum / n;
+
+    // Compute variance
+    let variance: f32 = x.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / n;
+    let inv_std = 1.0 / (variance + eps).sqrt();
+
+    // Apply normalization, weight, and bias
+    for i in 0..len {
+        let normalized = (x[i] - mean) * inv_std;
+        x[i] = normalized * weight[i] + bias[i];
+    }
+}
+
+/// Batched RMSNorm - process multiple vectors
+///
+/// # Arguments
+/// * `x` - Input tensor (batch_size, dim), modified in-place
+/// * `weight` - Shared weight parameters (dim,)
+/// * `batch_size` - Number of vectors in batch
+/// * `dim` - Dimension of each vector
+/// * `eps` - Numerical stability constant
+pub fn batched_rms_norm_neon(
+    x: &mut [f32],
+    weight: &[f32],
+    batch_size: usize,
+    dim: usize,
+    eps: f32,
+) {
+    debug_assert_eq!(x.len(), batch_size * dim);
+    debug_assert_eq!(weight.len(), dim);
+
+    for b in 0..batch_size {
+        let offset = b * dim;
+        rms_norm_neon(&mut x[offset..offset + dim], weight, eps);
+    }
+}
+
+/// Batched LayerNorm - process multiple vectors
+///
+/// # Arguments
+/// * `x` - Input tensor (batch_size, dim), modified in-place
+/// * `weight` - Shared gamma parameters (dim,)
+/// * `bias` - Shared beta parameters (dim,)
+/// * `batch_size` - Number of vectors in batch
+/// * `dim` - Dimension of each vector
+/// * `eps` - Numerical stability constant
+pub fn batched_layer_norm_neon(
+    x: &mut [f32],
+    weight: &[f32],
+    bias: &[f32],
+    batch_size: usize,
+    dim: usize,
+    eps: f32,
+) {
+    debug_assert_eq!(x.len(), batch_size * dim);
+    debug_assert_eq!(weight.len(), dim);
+    debug_assert_eq!(bias.len(), dim);
+
+    for b in 0..batch_size {
+        let offset = b * dim;
+        layer_norm_neon(&mut x[offset..offset + dim], weight, bias, eps);
+    }
+}
+
+/// Compute only the RMS value without applying normalization
+///
+/// Useful for monitoring activation magnitudes.
+#[inline(always)]
+pub fn compute_rms(x: &[f32]) -> f32 {
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        compute_rms_neon_impl(x)
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        compute_rms_scalar(x)
+    }
+}
+
+#[cfg(target_arch = "aarch64")]
+#[inline(always)]
+unsafe fn compute_rms_neon_impl(x: &[f32]) -> f32 {
+    let len = x.len();
+    if len == 0 {
+        return 0.0;
+    }
+
+    let x_ptr = x.as_ptr();
+    let mut sum = vdupq_n_f32(0.0);
+
+    let chunks = len / NEON_LANE_WIDTH;
+    let mut idx = 0usize;
+
+    for _ in 0..chunks {
+        let v = vld1q_f32(x_ptr.add(idx));
+        sum = vfmaq_f32(sum, v, v);
+        idx += 4;
+    }
+
+    let mut sum_sq = vaddvq_f32(sum);
+
+    for i in idx..len {
+        let v = *x_ptr.add(i);
+        sum_sq += v * v;
+    }
+
+    (sum_sq / len as f32).sqrt()
+}
+
+#[allow(dead_code)]
+fn compute_rms_scalar(x: &[f32]) -> f32 {
+    let sum_sq: f32 = x.iter().map(|v| v * v).sum();
+    (sum_sq / x.len() as f32).sqrt()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_rms_norm_basic() {
+        let mut x = vec![1.0, 2.0, 3.0, 4.0];
+        let weight = vec![1.0; 4];
+        let eps = 1e-6;
+
+        rms_norm_neon(&mut x, &weight, eps);
+
+        // Check that output is normalized
+        let rms: f32 = (x.iter().map(|v| v * v).sum::<f32>() / 4.0).sqrt();
+        // After normalization, the RMS should be close to 1
+        // (not exactly 1 because original values weren't unit RMS)
+        assert!(x.iter().all(|v| v.is_finite()));
+    }
+
+    #[test]
+    fn test_rms_norm_with_weight() {
+        let mut x = vec![1.0, 1.0, 1.0, 1.0];
+        let weight = vec![2.0, 2.0, 2.0, 2.0];
+        let eps = 1e-6;
+
+        rms_norm_neon(&mut x, &weight, eps);
+
+        // All equal inputs with equal weights should give equal outputs
+        let first = x[0];
+        assert!(x.iter().all(|&v| (v - first).abs() < 1e-5));
+    }
+
+    #[test]
+    fn test_layer_norm_basic() {
+        let mut x = vec![1.0, 2.0, 3.0, 4.0];
+        let weight = vec![1.0; 4];
+        let bias = vec![0.0; 4];
+        let eps = 1e-6;
+
+        layer_norm_neon(&mut x, &weight, &bias, eps);
+
+        // Check that mean is approximately 0
+        let mean: f32 = x.iter().sum::<f32>() / 4.0;
+        assert!(mean.abs() < 1e-5, "Mean should be ~0, got {}", mean);
+
+        // Check that variance is approximately 1
+        let var: f32 = x.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / 4.0;
+        assert!(
+            (var - 1.0).abs() < 1e-4,
+            "Variance should be ~1, got {}",
+            var
+        );
+    }
+
+    #[test]
+    fn test_layer_norm_with_bias() {
+        let mut x = vec![0.0, 0.0, 0.0, 0.0];
+        let weight = vec![1.0; 4];
+        let bias = vec![5.0; 4];
+        let eps = 1e-6;
+
+        layer_norm_neon(&mut x, &weight, &bias, eps);
+
+        // With zero input and bias, output should be approximately bias
+        // (normalized zero is zero, zero * weight + bias = bias)
+        for v in &x {
+            assert!((v - 5.0).abs() < 1e-4, "Expected ~5.0, got {}", v);
+        }
+    }
+
+    #[test]
+    fn test_rms_norm_large() {
+        let dim = 256;
+        let mut x: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
+        let weight = vec![1.0; dim];
+        let eps = 1e-6;
+
+        rms_norm_neon(&mut x, &weight, eps);
+
+        assert!(x.iter().all(|v| v.is_finite()));
+    }
+
+    #[test]
+    fn test_layer_norm_large() {
+        let dim = 256;
+        let mut x: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
+        let weight = vec![1.0; dim];
+        let bias = vec![0.0; dim];
+        let eps = 1e-6;
+
+        layer_norm_neon(&mut x, &weight, &bias, eps);
+
+        // Verify normalized mean and variance
+        let mean: f32 = x.iter().sum::<f32>() / dim as f32;
+        assert!(mean.abs() < 1e-4, "Mean should be ~0, got {}", mean);
+    }
+
+    #[test]
+    fn test_batched_rms_norm() {
+        let batch_size = 4;
+        let dim = 16;
+        let mut x: Vec<f32> = (0..batch_size * dim).map(|i| (i as f32) * 0.1).collect();
+        let weight = vec![1.0; dim];
+
+        batched_rms_norm_neon(&mut x, &weight, batch_size, dim, 1e-6);
+
+        assert!(x.iter().all(|v| v.is_finite()));
+    }
+
+    #[test]
+    fn test_batched_layer_norm() {
+        let batch_size = 4;
+        let dim = 16;
+        let mut x: Vec<f32> = (0..batch_size * dim).map(|i| (i as f32) * 0.1).collect();
+        let weight = vec![1.0; dim];
+        let bias = vec![0.0; dim];
+
+        batched_layer_norm_neon(&mut x, &weight, &bias, batch_size, dim, 1e-6);
+
+        // Check each batch vector is normalized
+        for b in 0..batch_size {
+            let offset = b * dim;
+            let slice = &x[offset..offset + dim];
+            let mean: f32 = slice.iter().sum::<f32>() / dim as f32;
+            assert!(
+                mean.abs() < 1e-4,
+                "Batch {} mean should be ~0, got {}",
+                b,
+                mean
+            );
+        }
+    }
+
+    #[test]
+    fn test_compute_rms() {
+        let x = vec![3.0, 4.0]; // RMS = sqrt((9+16)/2) = sqrt(12.5) ~ 3.536
+        let rms = compute_rms(&x);
+        assert!(
+            (rms - 3.5355).abs() < 0.01,
+            "RMS should be ~3.536, got {}",
+            rms
+        );
+    }
+
+    #[test]
+    fn test_rms_norm_matches_scalar() {
+        let dim = 64;
+        let mut x_neon: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.1 - 3.0).collect();
+        let mut x_scalar = x_neon.clone();
+        let weight: Vec<f32> = (0..dim).map(|i| 0.5 + (i as f32) * 0.01).collect();
+        let eps = 1e-6;
+
+        rms_norm_neon(&mut x_neon, &weight, eps);
+        rms_norm_scalar(&mut x_scalar, &weight, eps);
+
+        for i in 0..dim {
+            assert!(
+                (x_neon[i] - x_scalar[i]).abs() < 1e-4,
+                "Mismatch at {}: {} vs {}",
+                i,
+                x_neon[i],
+                x_scalar[i]
+            );
+        }
+    }
+
+    #[test]
+    fn test_layer_norm_matches_scalar() {
+        let dim = 64;
+        let mut x_neon: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.1 - 3.0).collect();
+        let mut x_scalar = x_neon.clone();
+        let weight: Vec<f32> = (0..dim).map(|i| 0.5 + (i as f32) * 0.01).collect();
+        let bias: Vec<f32> = (0..dim).map(|i| -0.2 + (i as f32) * 0.005).collect();
+        let eps = 1e-6;
+
+        layer_norm_neon(&mut x_neon, &weight, &bias, eps);
+        layer_norm_scalar(&mut x_scalar, &weight, &bias, eps);
+
+        for i in 0..dim {
+            assert!(
+                (x_neon[i] - x_scalar[i]).abs() < 1e-4,
+                "Mismatch at {}: {} vs {}",
+                i,
+                x_neon[i],
+                x_scalar[i]
+            );
+        }
+    }
+}

vendor/ruvector/crates/ruvllm/src/kernels/rope.rs

@@ -0,0 +1,659 @@
+//! NEON-Optimized Rotary Position Embeddings (RoPE)
+//!
+//! Implements efficient RoPE operations for transformer models:
+//!
+//! - **Standard RoPE**: Original rotary embeddings (Llama, GPT-NeoX)
+//! - **Scaled RoPE**: Position interpolation for extended context
+//! - **YaRN**: Yet another RoPE extension for very long contexts
+//!
+//! ## Mathematical Background
+//!
+//! RoPE applies rotation to query and key vectors based on position:
+//! ```text
+//! x_rotated = x * cos(theta) + rotate_half(x) * sin(theta)
+//! where theta = position * base^(-2i/d)
+//! ```
+//!
+//! ## Performance
+//!
+//! | Model | Head Dim | M4 Pro Throughput |
+//! |-------|----------|-------------------|
+//! | Llama-2 | 128 | ~4.2 GB/s |
+//! | Mistral | 128 | ~4.2 GB/s |
+//! | Llama-3 | 128 | ~4.0 GB/s (higher base) |
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+use super::{NEON_LANE_WIDTH, UNROLL_FACTOR};
+use std::f32::consts::PI;
+
+/// RoPE configuration
+#[derive(Debug, Clone, Copy)]
+pub struct RopeConfig {
+    /// Base frequency (10000.0 for Llama, 1000000.0 for some models)
+    pub base: f32,
+    /// Head dimension
+    pub head_dim: usize,
+    /// Maximum sequence length for precomputation
+    pub max_seq_len: usize,
+    /// Scaling factor for position interpolation (1.0 = no scaling)
+    pub scaling_factor: f32,
+    /// Whether to use NTK-aware scaling
+    pub ntk_aware: bool,
+    /// Original maximum sequence length (for scaling)
+    pub original_max_len: usize,
+}
+
+impl Default for RopeConfig {
+    fn default() -> Self {
+        Self {
+            base: 10000.0,
+            head_dim: 128,
+            max_seq_len: 4096,
+            scaling_factor: 1.0,
+            ntk_aware: false,
+            original_max_len: 4096,
+        }
+    }
+}
+
+impl RopeConfig {
+    /// Create config for Llama-2 style models
+    pub fn llama2(head_dim: usize, max_seq_len: usize) -> Self {
+        Self {
+            base: 10000.0,
+            head_dim,
+            max_seq_len,
+            ..Default::default()
+        }
+    }
+
+    /// Create config for Llama-3 style models (higher base)
+    pub fn llama3(head_dim: usize, max_seq_len: usize) -> Self {
+        Self {
+            base: 500000.0,
+            head_dim,
+            max_seq_len,
+            ..Default::default()
+        }
+    }
+
+    /// Create config for Mistral style models
+    pub fn mistral(head_dim: usize, max_seq_len: usize) -> Self {
+        Self {
+            base: 10000.0,
+            head_dim,
+            max_seq_len,
+            ..Default::default()
+        }
+    }
+
+    /// Create config with position interpolation
+    pub fn with_scaling(mut self, scaling_factor: f32) -> Self {
+        self.scaling_factor = scaling_factor;
+        self
+    }
+
+    /// Enable NTK-aware scaling
+    pub fn with_ntk(mut self, original_max_len: usize) -> Self {
+        self.ntk_aware = true;
+        self.original_max_len = original_max_len;
+        self
+    }
+
+    /// Compute effective base with NTK scaling
+    pub fn effective_base(&self) -> f32 {
+        if self.ntk_aware && self.max_seq_len > self.original_max_len {
+            let scale = self.max_seq_len as f32 / self.original_max_len as f32;
+            self.base * scale.powf((self.head_dim as f32) / (self.head_dim as f32 - 2.0))
+        } else {
+            self.base
+        }
+    }
+}
+
+/// Precomputed sin/cos tables for RoPE
+#[derive(Debug, Clone)]
+pub struct RopeTables {
+    /// Cosine values (max_seq_len, head_dim/2)
+    pub cos: Vec<f32>,
+    /// Sine values (max_seq_len, head_dim/2)
+    pub sin: Vec<f32>,
+    /// Half of head dimension
+    pub half_dim: usize,
+    /// Maximum sequence length
+    pub max_seq_len: usize,
+}
+
+impl RopeTables {
+    /// Get cos/sin for a specific position
+    #[inline(always)]
+    pub fn get(&self, position: usize) -> (&[f32], &[f32]) {
+        let offset = position * self.half_dim;
+        (
+            &self.cos[offset..offset + self.half_dim],
+            &self.sin[offset..offset + self.half_dim],
+        )
+    }
+}
+
+/// Precompute sin/cos tables for RoPE
+///
+/// # Arguments
+/// * `max_seq_len` - Maximum sequence length
+/// * `head_dim` - Dimension per head
+/// * `base` - RoPE base frequency
+///
+/// # Returns
+/// Tuple of (cos_table, sin_table), each of shape (max_seq_len, head_dim/2)
+pub fn precompute_rope_tables(
+    max_seq_len: usize,
+    head_dim: usize,
+    base: f32,
+) -> (Vec<f32>, Vec<f32>) {
+    let half_dim = head_dim / 2;
+    let mut cos_table = vec![0.0; max_seq_len * half_dim];
+    let mut sin_table = vec![0.0; max_seq_len * half_dim];
+
+    // Compute inverse frequencies: 1 / (base^(2i/d))
+    let inv_freq: Vec<f32> = (0..half_dim)
+        .map(|i| 1.0 / base.powf((2 * i) as f32 / head_dim as f32))
+        .collect();
+
+    // Compute sin/cos for each position
+    for pos in 0..max_seq_len {
+        let offset = pos * half_dim;
+        for (i, &freq) in inv_freq.iter().enumerate() {
+            let theta = pos as f32 * freq;
+            cos_table[offset + i] = theta.cos();
+            sin_table[offset + i] = theta.sin();
+        }
+    }
+
+    (cos_table, sin_table)
+}
+
+/// Precompute RoPE tables with configuration
+pub fn precompute_rope_tables_with_config(config: &RopeConfig) -> RopeTables {
+    let base = config.effective_base();
+    let (cos, sin) = precompute_rope_tables(config.max_seq_len, config.head_dim, base);
+
+    // Apply scaling factor if needed
+    let (cos, sin) = if config.scaling_factor != 1.0 {
+        let half_dim = config.head_dim / 2;
+        let mut scaled_cos = vec![0.0; config.max_seq_len * half_dim];
+        let mut scaled_sin = vec![0.0; config.max_seq_len * half_dim];
+
+        for pos in 0..config.max_seq_len {
+            let scaled_pos = pos as f32 / config.scaling_factor;
+            let lower_pos = scaled_pos.floor() as usize;
+            let upper_pos = (lower_pos + 1).min(config.max_seq_len - 1);
+            let frac = scaled_pos - lower_pos as f32;
+
+            let offset = pos * half_dim;
+            let lower_offset = lower_pos * half_dim;
+            let upper_offset = upper_pos * half_dim;
+
+            for i in 0..half_dim {
+                // Linear interpolation
+                scaled_cos[offset + i] =
+                    cos[lower_offset + i] * (1.0 - frac) + cos[upper_offset + i] * frac;
+                scaled_sin[offset + i] =
+                    sin[lower_offset + i] * (1.0 - frac) + sin[upper_offset + i] * frac;
+            }
+        }
+
+        (scaled_cos, scaled_sin)
+    } else {
+        (cos, sin)
+    };
+
+    RopeTables {
+        cos,
+        sin,
+        half_dim: config.head_dim / 2,
+        max_seq_len: config.max_seq_len,
+    }
+}
+
+/// Apply RoPE to query and key tensors in-place with NEON optimization
+///
+/// # Arguments
+/// * `x` - Input tensor to rotate (modified in-place)
+/// * `positions` - Position indices for each token
+/// * `head_dim` - Dimension per head
+/// * `base` - RoPE base frequency
+///
+/// # Implementation Details
+/// Uses interleaved rotation: pairs (x0, x1), (x2, x3), ... are rotated together
+#[inline(always)]
+pub fn apply_rope_neon(x: &mut [f32], positions: &[usize], head_dim: usize, base: f32) {
+    let half_dim = head_dim / 2;
+    let num_tokens = positions.len();
+    let stride = head_dim;
+
+    debug_assert_eq!(x.len(), num_tokens * head_dim);
+
+    // Precompute inverse frequencies
+    let inv_freq: Vec<f32> = (0..half_dim)
+        .map(|i| 1.0 / base.powf((2 * i) as f32 / head_dim as f32))
+        .collect();
+
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        apply_rope_neon_impl(x, positions, &inv_freq, half_dim, stride);
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        apply_rope_scalar(x, positions, &inv_freq, half_dim, stride);
+    }
+}
+
+/// Apply RoPE with precomputed tables
+#[inline(always)]
+pub fn apply_rope_with_tables(x: &mut [f32], positions: &[usize], tables: &RopeTables) {
+    let half_dim = tables.half_dim;
+    let num_tokens = positions.len();
+    let head_dim = half_dim * 2;
+
+    debug_assert_eq!(x.len(), num_tokens * head_dim);
+
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        apply_rope_tables_neon_impl(x, positions, tables, half_dim);
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        apply_rope_tables_scalar(x, positions, tables, half_dim);
+    }
+}
+
+/// NEON implementation of RoPE
+#[cfg(target_arch = "aarch64")]
+#[inline(always)]
+unsafe fn apply_rope_neon_impl(
+    x: &mut [f32],
+    positions: &[usize],
+    inv_freq: &[f32],
+    half_dim: usize,
+    stride: usize,
+) {
+    let x_ptr = x.as_mut_ptr();
+    let inv_freq_ptr = inv_freq.as_ptr();
+
+    for (tok_idx, &pos) in positions.iter().enumerate() {
+        let tok_offset = tok_idx * stride;
+
+        // Process in chunks of 4 (2 pairs at a time)
+        let chunks = half_dim / (NEON_LANE_WIDTH / 2);
+
+        let mut freq_idx = 0usize;
+        for _ in 0..chunks {
+            // Load inverse frequencies
+            let freq0 = *inv_freq_ptr.add(freq_idx);
+            let freq1 = *inv_freq_ptr.add(freq_idx + 1);
+
+            // Compute theta = position * inv_freq
+            let theta0 = pos as f32 * freq0;
+            let theta1 = pos as f32 * freq1;
+
+            // Compute sin/cos
+            let cos0 = theta0.cos();
+            let sin0 = theta0.sin();
+            let cos1 = theta1.cos();
+            let sin1 = theta1.sin();
+
+            // Load x values (pairs)
+            let x_offset = tok_offset + freq_idx * 2;
+            let x0 = *x_ptr.add(x_offset);
+            let x1 = *x_ptr.add(x_offset + 1);
+            let x2 = *x_ptr.add(x_offset + 2);
+            let x3 = *x_ptr.add(x_offset + 3);
+
+            // Apply rotation: x_new = x * cos - x_rotated * sin
+            // For pair (x0, x1): rotated is (-x1, x0)
+            *x_ptr.add(x_offset) = x0 * cos0 - x1 * sin0;
+            *x_ptr.add(x_offset + 1) = x1 * cos0 + x0 * sin0;
+            *x_ptr.add(x_offset + 2) = x2 * cos1 - x3 * sin1;
+            *x_ptr.add(x_offset + 3) = x3 * cos1 + x2 * sin1;
+
+            freq_idx += 2;
+        }
+
+        // Handle remaining pairs
+        while freq_idx < half_dim {
+            let freq = *inv_freq_ptr.add(freq_idx);
+            let theta = pos as f32 * freq;
+            let cos_val = theta.cos();
+            let sin_val = theta.sin();
+
+            let x_offset = tok_offset + freq_idx * 2;
+            let x0 = *x_ptr.add(x_offset);
+            let x1 = *x_ptr.add(x_offset + 1);
+
+            *x_ptr.add(x_offset) = x0 * cos_val - x1 * sin_val;
+            *x_ptr.add(x_offset + 1) = x1 * cos_val + x0 * sin_val;
+
+            freq_idx += 1;
+        }
+    }
+}
+
+/// NEON implementation with precomputed tables
+#[cfg(target_arch = "aarch64")]
+#[inline(always)]
+unsafe fn apply_rope_tables_neon_impl(
+    x: &mut [f32],
+    positions: &[usize],
+    tables: &RopeTables,
+    half_dim: usize,
+) {
+    let x_ptr = x.as_mut_ptr();
+    let head_dim = half_dim * 2;
+
+    for (tok_idx, &pos) in positions.iter().enumerate() {
+        debug_assert!(pos < tables.max_seq_len);
+
+        let tok_offset = tok_idx * head_dim;
+        let table_offset = pos * half_dim;
+
+        let cos_ptr = tables.cos.as_ptr().add(table_offset);
+        let sin_ptr = tables.sin.as_ptr().add(table_offset);
+
+        // Process with 4x unrolling
+        let chunks = half_dim / UNROLL_FACTOR;
+
+        let mut freq_idx = 0usize;
+        for _ in 0..chunks {
+            // Load cos/sin vectors
+            let cos_vec = vld1q_f32(cos_ptr.add(freq_idx));
+            let sin_vec = vld1q_f32(sin_ptr.add(freq_idx));
+
+            // Load x pairs (interleaved)
+            let x_offset = tok_offset + freq_idx * 2;
+
+            // Load 8 values (4 pairs)
+            let x_01 = vld1q_f32(x_ptr.add(x_offset));
+            let x_23 = vld1q_f32(x_ptr.add(x_offset + 4));
+
+            // Deinterleave to get even/odd elements
+            let x_even = vuzp1q_f32(x_01, x_23);
+            let x_odd = vuzp2q_f32(x_01, x_23);
+
+            // Apply rotation
+            // x_new_even = x_even * cos - x_odd * sin
+            // x_new_odd = x_odd * cos + x_even * sin
+            let x_new_even = vfmsq_f32(vmulq_f32(x_even, cos_vec), x_odd, sin_vec);
+            let x_new_odd = vfmaq_f32(vmulq_f32(x_odd, cos_vec), x_even, sin_vec);
+
+            // Interleave back
+            let out_01 = vzip1q_f32(x_new_even, x_new_odd);
+            let out_23 = vzip2q_f32(x_new_even, x_new_odd);
+
+            vst1q_f32(x_ptr.add(x_offset), out_01);
+            vst1q_f32(x_ptr.add(x_offset + 4), out_23);
+
+            freq_idx += 4;
+        }
+
+        // Handle remaining pairs
+        while freq_idx < half_dim {
+            let cos_val = *cos_ptr.add(freq_idx);
+            let sin_val = *sin_ptr.add(freq_idx);
+
+            let x_offset = tok_offset + freq_idx * 2;
+            let x0 = *x_ptr.add(x_offset);
+            let x1 = *x_ptr.add(x_offset + 1);
+
+            *x_ptr.add(x_offset) = x0 * cos_val - x1 * sin_val;
+            *x_ptr.add(x_offset + 1) = x1 * cos_val + x0 * sin_val;
+
+            freq_idx += 1;
+        }
+    }
+}
+
+/// Scalar fallback for RoPE
+#[allow(dead_code)]
+fn apply_rope_scalar(
+    x: &mut [f32],
+    positions: &[usize],
+    inv_freq: &[f32],
+    half_dim: usize,
+    stride: usize,
+) {
+    for (tok_idx, &pos) in positions.iter().enumerate() {
+        let tok_offset = tok_idx * stride;
+
+        for (i, &freq) in inv_freq.iter().enumerate() {
+            let theta = pos as f32 * freq;
+            let cos_val = theta.cos();
+            let sin_val = theta.sin();
+
+            let x_offset = tok_offset + i * 2;
+            let x0 = x[x_offset];
+            let x1 = x[x_offset + 1];
+
+            x[x_offset] = x0 * cos_val - x1 * sin_val;
+            x[x_offset + 1] = x1 * cos_val + x0 * sin_val;
+        }
+    }
+}
+
+/// Scalar fallback with precomputed tables
+#[allow(dead_code)]
+fn apply_rope_tables_scalar(
+    x: &mut [f32],
+    positions: &[usize],
+    tables: &RopeTables,
+    half_dim: usize,
+) {
+    let head_dim = half_dim * 2;
+
+    for (tok_idx, &pos) in positions.iter().enumerate() {
+        let tok_offset = tok_idx * head_dim;
+        let (cos_slice, sin_slice) = tables.get(pos);
+
+        for i in 0..half_dim {
+            let cos_val = cos_slice[i];
+            let sin_val = sin_slice[i];
+
+            let x_offset = tok_offset + i * 2;
+            let x0 = x[x_offset];
+            let x1 = x[x_offset + 1];
+
+            x[x_offset] = x0 * cos_val - x1 * sin_val;
+            x[x_offset + 1] = x1 * cos_val + x0 * sin_val;
+        }
+    }
+}
+
+/// Compute RoPE frequencies for a given position
+#[inline(always)]
+pub fn compute_rope_freqs(position: usize, head_dim: usize, base: f32) -> Vec<f32> {
+    let half_dim = head_dim / 2;
+    (0..half_dim)
+        .map(|i| {
+            let freq = 1.0 / base.powf((2 * i) as f32 / head_dim as f32);
+            position as f32 * freq
+        })
+        .collect()
+}
+
+/// Apply inverse RoPE (for position un-embedding)
+pub fn apply_inverse_rope_neon(x: &mut [f32], positions: &[usize], head_dim: usize, base: f32) {
+    let half_dim = head_dim / 2;
+    let stride = head_dim;
+
+    // Inverse RoPE uses negative angles
+    let inv_freq: Vec<f32> = (0..half_dim)
+        .map(|i| -1.0 / base.powf((2 * i) as f32 / head_dim as f32))
+        .collect();
+
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        apply_rope_neon_impl(x, positions, &inv_freq, half_dim, stride);
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        apply_rope_scalar(x, positions, &inv_freq, half_dim, stride);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_precompute_tables() {
+        let (cos, sin) = precompute_rope_tables(128, 64, 10000.0);
+
+        // Check dimensions
+        assert_eq!(cos.len(), 128 * 32);
+        assert_eq!(sin.len(), 128 * 32);
+
+        // Position 0 should have cos = 1, sin = 0
+        for i in 0..32 {
+            assert!((cos[i] - 1.0).abs() < 1e-5, "cos[{}] = {}", i, cos[i]);
+            assert!(sin[i].abs() < 1e-5, "sin[{}] = {}", i, sin[i]);
+        }
+    }
+
+    #[test]
+    fn test_rope_config() {
+        let config = RopeConfig::llama2(128, 4096);
+        assert_eq!(config.base, 10000.0);
+        assert_eq!(config.effective_base(), 10000.0);
+
+        let scaled_config = RopeConfig::llama2(128, 8192).with_ntk(4096);
+        assert!(scaled_config.effective_base() > 10000.0);
+    }
+
+    #[test]
+    fn test_apply_rope_basic() {
+        let head_dim = 8;
+        let mut x: Vec<f32> = vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0];
+        let positions = vec![0usize];
+
+        apply_rope_neon(&mut x, &positions, head_dim, 10000.0);
+
+        // At position 0, rotation should be identity (cos=1, sin=0)
+        assert!((x[0] - 1.0).abs() < 1e-5);
+        assert!(x[1].abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_apply_rope_rotation() {
+        let head_dim = 4;
+        let mut x: Vec<f32> = vec![1.0, 0.0, 1.0, 0.0];
+        let positions = vec![1usize]; // Position 1 should rotate
+
+        let original = x.clone();
+        apply_rope_neon(&mut x, &positions, head_dim, 10000.0);
+
+        // Values should change for non-zero position
+        // The rotation should not be identity
+        assert!(
+            (x[0] - original[0]).abs() > 1e-6 || (x[1] - original[1]).abs() > 1e-6,
+            "RoPE should rotate at position 1"
+        );
+    }
+
+    #[test]
+    fn test_rope_tables() {
+        let config = RopeConfig {
+            head_dim: 16,
+            max_seq_len: 32,
+            base: 10000.0,
+            ..Default::default()
+        };
+
+        let tables = precompute_rope_tables_with_config(&config);
+        assert_eq!(tables.half_dim, 8);
+        assert_eq!(tables.max_seq_len, 32);
+
+        let (cos0, sin0) = tables.get(0);
+        assert_eq!(cos0.len(), 8);
+        assert_eq!(sin0.len(), 8);
+    }
+
+    #[test]
+    fn test_apply_rope_with_tables() {
+        let config = RopeConfig {
+            head_dim: 8,
+            max_seq_len: 16,
+            base: 10000.0,
+            ..Default::default()
+        };
+
+        let tables = precompute_rope_tables_with_config(&config);
+
+        let mut x1: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
+        let mut x2 = x1.clone();
+        let positions = vec![5usize];
+
+        apply_rope_neon(&mut x1, &positions, config.head_dim, config.base);
+        apply_rope_with_tables(&mut x2, &positions, &tables);
+
+        // Both methods should produce same result
+        for i in 0..8 {
+            assert!(
+                (x1[i] - x2[i]).abs() < 1e-4,
+                "Mismatch at {}: {} vs {}",
+                i,
+                x1[i],
+                x2[i]
+            );
+        }
+    }
+
+    #[test]
+    fn test_inverse_rope() {
+        let head_dim = 8;
+        let mut x: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
+        let original = x.clone();
+        let positions = vec![5usize];
+
+        // Apply RoPE then inverse RoPE
+        apply_rope_neon(&mut x, &positions, head_dim, 10000.0);
+        apply_inverse_rope_neon(&mut x, &positions, head_dim, 10000.0);
+
+        // Should return to original
+        for i in 0..8 {
+            assert!(
+                (x[i] - original[i]).abs() < 1e-4,
+                "Inverse RoPE failed at {}: {} vs {}",
+                i,
+                x[i],
+                original[i]
+            );
+        }
+    }
+
+    #[test]
+    fn test_multiple_tokens() {
+        let head_dim = 4;
+        let mut x: Vec<f32> = vec![
+            1.0, 0.0, 1.0, 0.0, // Token 0
+            1.0, 0.0, 1.0, 0.0, // Token 1
+            1.0, 0.0, 1.0, 0.0, // Token 2
+        ];
+        let positions = vec![0usize, 1, 2];
+
+        apply_rope_neon(&mut x, &positions, head_dim, 10000.0);
+
+        // Token 0 should be unchanged (position 0)
+        assert!((x[0] - 1.0).abs() < 1e-5);
+
+        // Tokens 1 and 2 should be rotated
+        // Just verify they're different from original
+        assert!(x
+            .iter()
+            .skip(4)
+            .any(|&v| (v - 1.0).abs() > 1e-5 || v.abs() > 1e-5));
+    }
+}