diff --git a/CHANGELOG.md b/CHANGELOG.md index 04bc65f53..0d81a5620 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,21 @@ ## [Unreleased] +### Added + +- **Native Q5_0 / Q5_1 packed matmul kernels (FFM, Kotlin/Native, JNI).** 0.39.0 shipped + packed GGUF *loading* for Q5_0/Q5_1 plus scalar + Panama kernels, but the native tier had + no Q5_x kernels — on the JVM the registry cascaded to Panama (50), and on Kotlin/Native and + Android the formats ran on the priority-0 scalar floor. New `skainet_q5_0_matmul` / + `skainet_q5_1_matmul` C kernels (plain NEON, no dotprod/i8mm requirement — runs on every + AArch64 core) expand the `qh` high-bit plane with a per-lane `vtstq_u8` bit test and fold + the dequant algebraically (`d*(dot - 16*Σx)` for Q5_0, `d*dot + m*Σx` for Q5_1) so the + per-block input sum hoists out of the output-row loop. Wired into all three consumers: + the FFM `NativeKernelProvider` (JVM), the cinterop `NativeKnKernelProvider` + (Kotlin/Native), and the Android JNI bridge (`JniKernels.q50Matmul`/`q51Matmul` + + `JniKernelProvider`), each with parity tests against the scalar references. Unblocks the + packed Q5_1 path for `functiongemma-270m` "Q5_K_M" checkpoints (whose attention/FFN + weights are Q5_1) under `NATIVE_OPTIMIZED` — see SKaiNET-transformers#170. (#708) ### Performance - **Primitive FP32 fast paths for the eager CPU ops** (`skainet-backend-cpu`, diff --git a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt index 174f3fafe..002a64f8c 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-jni-cpu/native/CMakeLists.txt @@ -11,6 +11,8 @@ set(SKAINET_KERNEL_SOURCES ${SKAINET_KERNELS_ROOT}/src/bf16_matmul.c ${SKAINET_KERNELS_ROOT}/src/fp16_matmul.c ${SKAINET_KERNELS_ROOT}/src/q4_0_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q5_0_matmul.c + ${SKAINET_KERNELS_ROOT}/src/q5_1_matmul.c ${SKAINET_KERNELS_ROOT}/src/q8_0_matmul.c ${SKAINET_KERNELS_ROOT}/src/q4k_matmul.c ${SKAINET_KERNELS_ROOT}/src/q5k_matmul.c diff --git a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c index 84bcc8cac..9d3a0bce2 100644 --- a/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c +++ b/skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c @@ -73,6 +73,34 @@ Java_sk_ainet_exec_kernel_jni_JniKernels_q40Matmul( inputDim, outputDim, out, outputOffset)) } +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q50Matmul( + JNIEnv* env, jobject thiz, + jfloatArray input, jint inputOffset, + jbyteArray weight, jint weightByteOffset, + jint inputDim, jint outputDim, + jfloatArray output, jint outputOffset +) { + (void) thiz; + SKAINET_JNI_MATMUL_BODY( + skainet_q5_0_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + +JNIEXPORT void JNICALL +Java_sk_ainet_exec_kernel_jni_JniKernels_q51Matmul( + JNIEnv* env, jobject thiz, + jfloatArray input, jint inputOffset, + jbyteArray weight, jint weightByteOffset, + jint inputDim, jint outputDim, + jfloatArray output, jint outputOffset +) { + (void) thiz; + SKAINET_JNI_MATMUL_BODY( + skainet_q5_1_matmul(in, inputOffset, (const uint8_t*) w, weightByteOffset, + inputDim, outputDim, out, outputOffset)) +} + JNIEXPORT void JNICALL Java_sk_ainet_exec_kernel_jni_JniKernels_q4kMatmul( JNIEnv* env, jobject thiz, diff --git a/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt b/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt index bcc51ffad..355559bff 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt +++ b/skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt @@ -8,6 +8,8 @@ import org.junit.Test import org.junit.runner.RunWith import sk.ainet.exec.kernel.ScalarQ4_0MatmulKernel import sk.ainet.exec.kernel.ScalarQ4_KMatmulKernel +import sk.ainet.exec.kernel.ScalarQ5_0MatmulKernel +import sk.ainet.exec.kernel.ScalarQ5_1MatmulKernel import sk.ainet.exec.kernel.ScalarQ6_KMatmulKernel import sk.ainet.exec.kernel.ScalarQ8_0MatmulKernel import kotlin.math.abs @@ -100,6 +102,8 @@ class JniKernelParityTest { private val conditionQ8_0: (ByteArray, Int) -> Unit = { b, base -> fp16One(b, base) } private val conditionQ4_0: (ByteArray, Int) -> Unit = { b, base -> fp16One(b, base) } + private val conditionQ5_0: (ByteArray, Int) -> Unit = { b, base -> fp16One(b, base) } + private val conditionQ5_1: (ByteArray, Int) -> Unit = { b, base -> fp16One(b, base); fp16One(b, base + 2) } private val conditionQ4K: (ByteArray, Int) -> Unit = { b, base -> fp16One(b, base); fp16One(b, base + 2) } private val conditionQ6K: (ByteArray, Int) -> Unit = { b, base -> fp16One(b, base + 208) } @@ -181,6 +185,18 @@ class JniKernelParityTest { reference = ScalarQ4_0MatmulKernel::matmul, jni = JniKernels::q40Matmul, ) + @Test + fun q50_parity() = assertExactParity( + 1024, 64, 11, 2e-1f, 32, 22, conditionQ5_0, + reference = ScalarQ5_0MatmulKernel::matmul, jni = JniKernels::q50Matmul, + ) + + @Test + fun q51_parity() = assertExactParity( + 1024, 64, 13, 2e-1f, 32, 24, conditionQ5_1, + reference = ScalarQ5_1MatmulKernel::matmul, jni = JniKernels::q51Matmul, + ) + @Test fun q4k_parity_single_block() = assertRmsParity( 256, 16, 42, 1e-2f, 256, 144, conditionQ4K, diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt index 10781bd02..fc77963d1 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernelProvider.kt @@ -4,6 +4,8 @@ import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.Q4KMatmulKernel import sk.ainet.backend.api.kernel.Q4_0MatmulKernel import sk.ainet.backend.api.kernel.Q5KMatmulKernel +import sk.ainet.backend.api.kernel.Q5_0MatmulKernel +import sk.ainet.backend.api.kernel.Q5_1MatmulKernel import sk.ainet.backend.api.kernel.Q6KMatmulKernel import sk.ainet.backend.api.kernel.Q8_0MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel @@ -55,6 +57,10 @@ public object JniKernelProvider : KernelProvider { override fun matmulQ6K(): Q6KMatmulKernel? = if (available) JniQ6KMatmul else null + override fun matmulQ5_0(): Q5_0MatmulKernel? = if (available) JniQ5_0Matmul else null + + override fun matmulQ5_1(): Q5_1MatmulKernel? = if (available) JniQ5_1Matmul else null + private object JniQ8_0Matmul : Q8_0MatmulKernel { override fun matmul( input: FloatArray, inputOffset: Int, @@ -77,6 +83,28 @@ public object JniKernelProvider : KernelProvider { ) } + private object JniQ5_0Matmul : Q5_0MatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q50Matmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } + + private object JniQ5_1Matmul : Q5_1MatmulKernel { + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ): Unit = JniKernels.q51Matmul( + input, inputOffset, weight, weightByteOffset, inputDim, outputDim, output, outputOffset + ) + } + private object JniQ4KMatmul : Q4KMatmulKernel { override fun matmul( input: FloatArray, inputOffset: Int, diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt index e2ee855bd..465b98b9b 100644 --- a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt @@ -101,6 +101,20 @@ public object JniKernels { output: FloatArray, outputOffset: Int, ) + public external fun q50Matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) + + public external fun q51Matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) + public external fun q4kMatmul( input: FloatArray, inputOffset: Int, weight: ByteArray, weightByteOffset: Int, diff --git a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt index 88d3bfd3a..447375b38 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt @@ -19,6 +19,8 @@ set(SKAINET_KERNEL_SOURCES src/fp16_matmul.c src/q8_0_matmul.c src/q4_0_matmul.c + src/q5_0_matmul.c + src/q5_1_matmul.c ) # SHARED: consumed by the JVM via java.lang.foreign (FFM), bundled as a JAR diff --git a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h index 54424ea8e..e3e960a63 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h +++ b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h @@ -214,6 +214,62 @@ SKAINET_API void skainet_q4_0_matmul( int32_t output_offset ); +/* + * Q5_0 matrix-vector multiply. + * + * output[output_offset + o] = sum_j input[input_offset + j] * + * dequant(weight[block, o, j]) + * + * Block layout: canonical ggml Q5_0, 32 elements per block, 22 bytes + * per block (2 B FP16 scale `d` + 4 B `qh` high-bit plane + 16 B packed + * 4-bit codes in split layout — low nibbles → elements 0..15, high + * nibbles → 16..31; bit j of the little-endian 32-bit `qh` is the fifth + * bit of element j), with packed weights laid out as + * weight + weight_byte_offset + (block_idx * output_dim + o) * 22 + * + * Dequant per element: `(code - 16) * d` with + * `code = nibble | (fifth_bit << 4)`. input_dim must be a multiple + * of 32. + */ +SKAINET_API void skainet_q5_0_matmul( + const float* input, + int32_t input_offset, + const uint8_t* weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* output, + int32_t output_offset +); + +/* + * Q5_1 matrix-vector multiply. + * + * output[output_offset + o] = sum_j input[input_offset + j] * + * dequant(weight[block, o, j]) + * + * Block layout: canonical ggml Q5_1, 32 elements per block, 24 bytes + * per block (2 B FP16 scale `d` + 2 B FP16 min `m` + 4 B `qh` high-bit + * plane + 16 B packed 4-bit codes in split layout; bit j of the + * little-endian 32-bit `qh` is the fifth bit of element j), with packed + * weights laid out as + * weight + weight_byte_offset + (block_idx * output_dim + o) * 24 + * + * Dequant per element: `d * code + m` (affine, no re-centring) with + * `code = nibble | (fifth_bit << 4)`. input_dim must be a multiple + * of 32. + */ +SKAINET_API void skainet_q5_1_matmul( + const float* input, + int32_t input_offset, + const uint8_t* weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* output, + int32_t output_offset +); + #ifdef __cplusplus } #endif diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/q5_0_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/q5_0_matmul.c new file mode 100644 index 000000000..8851e904f --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/q5_0_matmul.c @@ -0,0 +1,164 @@ +#include "skainet_kernels.h" +#include "skainet_simd.h" + +#include +#include +#include + +/* + * Native FP32 × Q5_0 matrix-vector matmul matching the + * sk.ainet.backend.api.kernel.Q5_0MatmulKernel SPI. + * + * Block layout (canonical ggml Q5_0, 32 elements, 22 bytes): + * - bytes 0..1 : FP16 little-endian scale `d` + * - bytes 2..5 : `qh` — 32-bit little-endian high-bit plane, bit j is + * the fifth bit of element j + * - bytes 6..21 : 16 bytes packing 32 4-bit codes in the *split* + * layout — low nibbles decode elements 0..15, high nibbles decode + * elements 16..31. + * + * Per-block packed weight layout: + * weight + weight_byte_offset + (block_idx * output_dim + o) * 22 + * + * Dequant per element: `(code - 16) * d` with + * `code = nibble | (fifth_bit << 4)` (unsigned, 0..31). The kernel uses + * the algebraic split + * + * sum_j x_j * d * (code_j - 16) = d * (dot(x, code) - 16 * sum(x)) + * + * so the inner loop accumulates the *unsigned* code dot product and the + * per-block input sum is hoisted OUT of the output-row loop (it only + * depends on the activations, not on `o`). Scale folding happens once + * per block, exactly like q4_0/q8_0. + * + * Loop order: block OUTER, output row INNER — see q8_0_matmul.c for the + * rationale (sequential weight reads; per-row accumulation order stays + * ascending-block, so results don't depend on output_dim). + * + * NEON path (SKAINET_HAVE_NEON): plain NEON only — no dotprod/i8mm + * requirement, so the body runs on every AArch64 core. The high-bit + * plane is expanded with a per-lane vtstq_u8 against a {1,2,4,…,128} + * bitmask after broadcasting each `qh` byte to its 8 lanes, masked to + * the fifth-bit value (16) and OR'd onto the nibbles; the resulting + * unsigned 5-bit codes widen to f32 FMA lanes via + * skainet_neon_u8x16_to_f32x4x4, the same structure as q8_0/q4_0. + */ + +/* Portable FP16 → FP32 conversion. Matches the Kotlin `decodeHalf` + * algorithm bit-for-bit. */ +static inline float skainet_q5_0_fp16_to_fp32(uint16_t h) { + uint32_t sign = ((uint32_t)(h & 0x8000u)) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t mant = h & 0x3FFu; + uint32_t bits; + if (exp == 0) { + if (mant == 0) { + bits = sign; + } else { + int e = -14; + while ((mant & 0x400u) == 0) { + mant <<= 1; + --e; + } + mant &= 0x3FFu; + bits = sign | ((uint32_t)(e + 127) << 23) | (mant << 13); + } + } else if (exp == 0x1Fu) { + bits = sign | 0x7F800000u | (mant << 13); + } else { + bits = sign | ((uint32_t)(exp - 15 + 127) << 23) | (mant << 13); + } + float r; + memcpy(&r, &bits, sizeof(r)); + return r; +} + +SKAINET_API void skainet_q5_0_matmul( + const float* SKAINET_RESTRICT input, int32_t input_offset, + const uint8_t* SKAINET_RESTRICT weight, int32_t weight_byte_offset, + int32_t input_dim, int32_t output_dim, + float* SKAINET_RESTRICT output, int32_t output_offset +) { + if (output_dim <= 0) return; + if (input_dim <= 0) { + for (int32_t o = 0; o < output_dim; ++o) { + output[output_offset + o] = 0.0f; + } + return; + } + + const int32_t BLOCK_SIZE = 32; + const int32_t BYTES_PER_BLOCK = 22; + const int32_t blocks_per_input_dim = input_dim / BLOCK_SIZE; + float* SKAINET_RESTRICT out_base = output + output_offset; + + for (int32_t o = 0; o < output_dim; ++o) out_base[o] = 0.0f; + + for (int32_t block_idx = 0; block_idx < blocks_per_input_dim; ++block_idx) { + const float* SKAINET_RESTRICT input_block = + input + input_offset + (size_t) block_idx * BLOCK_SIZE; + const uint8_t* SKAINET_RESTRICT block = + weight + weight_byte_offset + + (size_t)(block_idx * output_dim) * BYTES_PER_BLOCK; + + /* Depends only on the activations — hoisted out of the o-loop. */ + float input_sum = 0.0f; + for (int32_t k = 0; k < BLOCK_SIZE; ++k) input_sum += input_block[k]; + + for (int32_t o = 0; o < output_dim; ++o, block += BYTES_PER_BLOCK) { + uint16_t d_bits = (uint16_t) block[0] | ((uint16_t) block[1] << 8); + float d = skainet_q5_0_fp16_to_fp32(d_bits); + const uint8_t* SKAINET_RESTRICT qh = block + 2; + const uint8_t* SKAINET_RESTRICT qs = block + 6; + float code_dot = 0.0f; +#ifdef SKAINET_HAVE_NEON + /* Broadcast each qh byte to its 8 lanes, test the per-lane bit, + * mask to the fifth-bit value 16 and OR onto the nibbles. */ + const uint8x16_t bitmask = vcombine_u8( + vcreate_u8(0x8040201008040201ULL), + vcreate_u8(0x8040201008040201ULL)); + const uint8x8_t qh_bytes = vreinterpret_u8_u32(vdup_n_u32( + (uint32_t) qh[0] | ((uint32_t) qh[1] << 8) | + ((uint32_t) qh[2] << 16) | ((uint32_t) qh[3] << 24))); + const uint8x16_t qh_lo = vcombine_u8( + vdup_lane_u8(qh_bytes, 0), vdup_lane_u8(qh_bytes, 1)); + const uint8x16_t qh_hi = vcombine_u8( + vdup_lane_u8(qh_bytes, 2), vdup_lane_u8(qh_bytes, 3)); + const uint8x16_t fifth = vdupq_n_u8(0x10); + const uint8x16_t fifth_lo = vandq_u8(vtstq_u8(qh_lo, bitmask), fifth); + const uint8x16_t fifth_hi = vandq_u8(vtstq_u8(qh_hi, bitmask), fifth); + + const uint8x16_t packed = vld1q_u8(qs); + const uint8x16_t code_lo = vorrq_u8( + vandq_u8(packed, vdupq_n_u8(0x0F)), fifth_lo); /* elems 0..15 */ + const uint8x16_t code_hi = vorrq_u8( + vshrq_n_u8(packed, 4), fifth_hi); /* elems 16..31 */ + + float32x4_t lo_f[4]; + float32x4_t hi_f[4]; + skainet_neon_u8x16_to_f32x4x4(code_lo, lo_f); + skainet_neon_u8x16_to_f32x4x4(code_hi, hi_f); + float32x4_t accv = vdupq_n_f32(0.0f); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 0), lo_f[0]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 4), lo_f[1]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 8), lo_f[2]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 12), lo_f[3]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 16), hi_f[0]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 20), hi_f[1]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 24), hi_f[2]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 28), hi_f[3]); + code_dot = skainet_neon_hadd_f32(accv); +#else + uint32_t qh32 = (uint32_t) qh[0] | ((uint32_t) qh[1] << 8) | + ((uint32_t) qh[2] << 16) | ((uint32_t) qh[3] << 24); + for (int32_t k = 0; k < 16; ++k) { + int32_t lo = (int32_t)(qs[k] & 0x0F) | (int32_t)(((qh32 >> k) & 1u) << 4); + int32_t hi = (int32_t)(qs[k] >> 4) | (int32_t)(((qh32 >> (k + 16)) & 1u) << 4); + code_dot += input_block[k] * (float) lo; + code_dot += input_block[k + 16] * (float) hi; + } +#endif + out_base[o] += d * (code_dot - 16.0f * input_sum); + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/q5_1_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/q5_1_matmul.c new file mode 100644 index 000000000..d47842c2d --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/q5_1_matmul.c @@ -0,0 +1,164 @@ +#include "skainet_kernels.h" +#include "skainet_simd.h" + +#include +#include +#include + +/* + * Native FP32 × Q5_1 matrix-vector matmul matching the + * sk.ainet.backend.api.kernel.Q5_1MatmulKernel SPI. + * + * Block layout (canonical ggml Q5_1, 32 elements, 24 bytes): + * - bytes 0..1 : FP16 little-endian scale `d` + * - bytes 2..3 : FP16 little-endian min `m` + * - bytes 4..7 : `qh` — 32-bit little-endian high-bit plane, bit j is + * the fifth bit of element j + * - bytes 8..23 : 16 bytes packing 32 4-bit codes in the *split* + * layout — low nibbles decode elements 0..15, high nibbles decode + * elements 16..31. + * + * Per-block packed weight layout: + * weight + weight_byte_offset + (block_idx * output_dim + o) * 24 + * + * Dequant per element: `d * code + m` with + * `code = nibble | (fifth_bit << 4)` (unsigned, 0..31, affine — no + * re-centring). The kernel uses the algebraic split + * + * sum_j x_j * (d * code_j + m) = d * dot(x, code) + m * sum(x) + * + * so the inner loop accumulates the unsigned code dot product and the + * per-block input sum is hoisted OUT of the output-row loop (it only + * depends on the activations, not on `o`). `d`/`m` fold once per block. + * + * Loop order: block OUTER, output row INNER — see q8_0_matmul.c for the + * rationale (sequential weight reads; per-row accumulation order stays + * ascending-block, so results don't depend on output_dim). + * + * NEON path (SKAINET_HAVE_NEON): plain NEON only — no dotprod/i8mm + * requirement, so the body runs on every AArch64 core. Identical + * high-bit expansion + unsigned widen + vfmaq_f32 structure as + * q5_0_matmul.c; only the per-block fold differs (`+ m*sum` instead of + * `- 16d*sum`). + */ + +/* Portable FP16 → FP32 conversion. Matches the Kotlin `decodeHalf` + * algorithm bit-for-bit. */ +static inline float skainet_q5_1_fp16_to_fp32(uint16_t h) { + uint32_t sign = ((uint32_t)(h & 0x8000u)) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t mant = h & 0x3FFu; + uint32_t bits; + if (exp == 0) { + if (mant == 0) { + bits = sign; + } else { + int e = -14; + while ((mant & 0x400u) == 0) { + mant <<= 1; + --e; + } + mant &= 0x3FFu; + bits = sign | ((uint32_t)(e + 127) << 23) | (mant << 13); + } + } else if (exp == 0x1Fu) { + bits = sign | 0x7F800000u | (mant << 13); + } else { + bits = sign | ((uint32_t)(exp - 15 + 127) << 23) | (mant << 13); + } + float r; + memcpy(&r, &bits, sizeof(r)); + return r; +} + +SKAINET_API void skainet_q5_1_matmul( + const float* SKAINET_RESTRICT input, int32_t input_offset, + const uint8_t* SKAINET_RESTRICT weight, int32_t weight_byte_offset, + int32_t input_dim, int32_t output_dim, + float* SKAINET_RESTRICT output, int32_t output_offset +) { + if (output_dim <= 0) return; + if (input_dim <= 0) { + for (int32_t o = 0; o < output_dim; ++o) { + output[output_offset + o] = 0.0f; + } + return; + } + + const int32_t BLOCK_SIZE = 32; + const int32_t BYTES_PER_BLOCK = 24; + const int32_t blocks_per_input_dim = input_dim / BLOCK_SIZE; + float* SKAINET_RESTRICT out_base = output + output_offset; + + for (int32_t o = 0; o < output_dim; ++o) out_base[o] = 0.0f; + + for (int32_t block_idx = 0; block_idx < blocks_per_input_dim; ++block_idx) { + const float* SKAINET_RESTRICT input_block = + input + input_offset + (size_t) block_idx * BLOCK_SIZE; + const uint8_t* SKAINET_RESTRICT block = + weight + weight_byte_offset + + (size_t)(block_idx * output_dim) * BYTES_PER_BLOCK; + + /* Depends only on the activations — hoisted out of the o-loop. */ + float input_sum = 0.0f; + for (int32_t k = 0; k < BLOCK_SIZE; ++k) input_sum += input_block[k]; + + for (int32_t o = 0; o < output_dim; ++o, block += BYTES_PER_BLOCK) { + uint16_t d_bits = (uint16_t) block[0] | ((uint16_t) block[1] << 8); + uint16_t m_bits = (uint16_t) block[2] | ((uint16_t) block[3] << 8); + float d = skainet_q5_1_fp16_to_fp32(d_bits); + float m = skainet_q5_1_fp16_to_fp32(m_bits); + const uint8_t* SKAINET_RESTRICT qh = block + 4; + const uint8_t* SKAINET_RESTRICT qs = block + 8; + float code_dot = 0.0f; +#ifdef SKAINET_HAVE_NEON + /* Broadcast each qh byte to its 8 lanes, test the per-lane bit, + * mask to the fifth-bit value 16 and OR onto the nibbles. */ + const uint8x16_t bitmask = vcombine_u8( + vcreate_u8(0x8040201008040201ULL), + vcreate_u8(0x8040201008040201ULL)); + const uint8x8_t qh_bytes = vreinterpret_u8_u32(vdup_n_u32( + (uint32_t) qh[0] | ((uint32_t) qh[1] << 8) | + ((uint32_t) qh[2] << 16) | ((uint32_t) qh[3] << 24))); + const uint8x16_t qh_lo = vcombine_u8( + vdup_lane_u8(qh_bytes, 0), vdup_lane_u8(qh_bytes, 1)); + const uint8x16_t qh_hi = vcombine_u8( + vdup_lane_u8(qh_bytes, 2), vdup_lane_u8(qh_bytes, 3)); + const uint8x16_t fifth = vdupq_n_u8(0x10); + const uint8x16_t fifth_lo = vandq_u8(vtstq_u8(qh_lo, bitmask), fifth); + const uint8x16_t fifth_hi = vandq_u8(vtstq_u8(qh_hi, bitmask), fifth); + + const uint8x16_t packed = vld1q_u8(qs); + const uint8x16_t code_lo = vorrq_u8( + vandq_u8(packed, vdupq_n_u8(0x0F)), fifth_lo); /* elems 0..15 */ + const uint8x16_t code_hi = vorrq_u8( + vshrq_n_u8(packed, 4), fifth_hi); /* elems 16..31 */ + + float32x4_t lo_f[4]; + float32x4_t hi_f[4]; + skainet_neon_u8x16_to_f32x4x4(code_lo, lo_f); + skainet_neon_u8x16_to_f32x4x4(code_hi, hi_f); + float32x4_t accv = vdupq_n_f32(0.0f); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 0), lo_f[0]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 4), lo_f[1]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 8), lo_f[2]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 12), lo_f[3]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 16), hi_f[0]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 20), hi_f[1]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 24), hi_f[2]); + accv = vfmaq_f32(accv, vld1q_f32(input_block + 28), hi_f[3]); + code_dot = skainet_neon_hadd_f32(accv); +#else + uint32_t qh32 = (uint32_t) qh[0] | ((uint32_t) qh[1] << 8) | + ((uint32_t) qh[2] << 16) | ((uint32_t) qh[3] << 24); + for (int32_t k = 0; k < 16; ++k) { + int32_t lo = (int32_t)(qs[k] & 0x0F) | (int32_t)(((qh32 >> k) & 1u) << 4); + int32_t hi = (int32_t)(qs[k] >> 4) | (int32_t)(((qh32 >> (k + 16)) & 1u) << 4); + code_dot += input_block[k] * (float) lo; + code_dot += input_block[k + 16] * (float) hi; + } +#endif + out_base[o] += d * code_dot + m * input_sum; + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt index d15b37a74..cff68465f 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt @@ -9,6 +9,8 @@ import sk.ainet.backend.api.kernel.Q4KMatmulKernel import sk.ainet.backend.api.kernel.Q4KMemSegMatmulKernel import sk.ainet.backend.api.kernel.Q4_0MatmulKernel import sk.ainet.backend.api.kernel.Q5KMatmulKernel +import sk.ainet.backend.api.kernel.Q5_0MatmulKernel +import sk.ainet.backend.api.kernel.Q5_1MatmulKernel import sk.ainet.backend.api.kernel.Q6KMatmulKernel import sk.ainet.backend.api.kernel.Q8_0MatmulKernel @@ -113,4 +115,10 @@ public object NativeKernelProvider : KernelProvider, MemSegKernelProvider { override fun matmulQ6K(): Q6KMatmulKernel? = if (NativeQ6KMatmulKernel.isAvailable()) NativeQ6KMatmulKernel else null + + override fun matmulQ5_0(): Q5_0MatmulKernel? = + if (NativeQ5_0MatmulKernel.isAvailable()) NativeQ5_0MatmulKernel else null + + override fun matmulQ5_1(): Q5_1MatmulKernel? = + if (NativeQ5_1MatmulKernel.isAvailable()) NativeQ5_1MatmulKernel else null } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt new file mode 100644 index 000000000..bb6da5038 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt @@ -0,0 +1,105 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.FunctionDescriptor +import java.lang.foreign.Linker +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.lang.invoke.MethodHandle +import sk.ainet.backend.api.kernel.Q5_0MatmulKernel + +/** + * Native (FFM) implementation of [Q5_0MatmulKernel]. + * + * Wraps the bundled C symbol + * + * void skainet_q5_0_matmul( + * const float* input, int32_t input_offset, + * const uint8_t* weight, int32_t weight_byte_offset, + * int32_t input_dim, int32_t output_dim, + * float* output, int32_t output_offset); + * + * The C kernel decodes the ggml-canonical Q5_0 block (FP16 scale + 4-byte + * `qh` high-bit plane + 16 packed bytes, split nibble layout) with + * `(code - 16) * d` dequant; on AArch64 a plain-NEON body (no dotprod + * requirement) expands the high-bit plane with `vtstq_u8` and FMAs the + * unsigned codes against the activation lanes. + * + * Numerical parity vs [ScalarQ5_0MatmulKernel] is asserted by + * `NativeQ5_0MatmulKernelParityTest` within the same `1e-2 * + * blocksPerInputDim` band the Panama parity uses. + */ +internal object NativeQ5_0MatmulKernel : Q5_0MatmulKernel { + + fun isAvailable(): Boolean = handle != null + + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) { + require(inputDim % BLOCK_SIZE == 0) { + "NativeQ5_0MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" + } + if (outputDim == 0) return + + val mh = handle + ?: error("NativeQ5_0MatmulKernel.matmul invoked while native library unavailable") + + val blocksPerInputDim = inputDim / BLOCK_SIZE + val inputReachFloats = if (inputDim == 0) 0 else inputOffset + inputDim + val weightReachBytes = if (inputDim == 0 || outputDim == 0) 0 + else weightByteOffset + blocksPerInputDim * outputDim * BYTES_PER_BLOCK + val outputReachFloats = outputOffset + outputDim + + Arena.ofConfined().use { arena -> + val fAlign = ValueLayout.JAVA_FLOAT.byteAlignment() + val bAlign = ValueLayout.JAVA_BYTE.byteAlignment() + + val inputSeg: MemorySegment = if (inputReachFloats > 0) + arena.allocate(inputReachFloats.toLong() * java.lang.Float.BYTES, fAlign) + else MemorySegment.NULL + val weightSeg: MemorySegment = if (weightReachBytes > 0) + arena.allocate(weightReachBytes.toLong(), bAlign) + else MemorySegment.NULL + val outputSeg: MemorySegment = + arena.allocate(outputReachFloats.toLong() * java.lang.Float.BYTES, fAlign) + + if (inputReachFloats > 0) { + MemorySegment.copy(input, 0, inputSeg, ValueLayout.JAVA_FLOAT, 0L, inputReachFloats) + } + if (weightReachBytes > 0) { + MemorySegment.copy(weight, 0, weightSeg, ValueLayout.JAVA_BYTE, 0L, weightReachBytes) + } + + mh.invoke( + inputSeg, inputOffset, + weightSeg, weightByteOffset, + inputDim, outputDim, + outputSeg, outputOffset, + ) + + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, outputReachFloats) + } + } + + private const val BLOCK_SIZE = 32 + private const val BYTES_PER_BLOCK = 22 + + private val handle: MethodHandle? by lazy { + val lookup = NativeLibraryLoader.lookup() ?: return@lazy null + val symbol = lookup.find("skainet_q5_0_matmul").orElse(null) ?: return@lazy null + val descriptor = FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, // input + ValueLayout.JAVA_INT, // input_offset + ValueLayout.ADDRESS, // weight + ValueLayout.JAVA_INT, // weight_byte_offset + ValueLayout.JAVA_INT, // input_dim + ValueLayout.JAVA_INT, // output_dim + ValueLayout.ADDRESS, // output + ValueLayout.JAVA_INT, // output_offset + ) + runCatching { Linker.nativeLinker().downcallHandle(symbol, descriptor) }.getOrNull() + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt new file mode 100644 index 000000000..69cbb897f --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt @@ -0,0 +1,105 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.FunctionDescriptor +import java.lang.foreign.Linker +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.lang.invoke.MethodHandle +import sk.ainet.backend.api.kernel.Q5_1MatmulKernel + +/** + * Native (FFM) implementation of [Q5_1MatmulKernel]. + * + * Wraps the bundled C symbol + * + * void skainet_q5_1_matmul( + * const float* input, int32_t input_offset, + * const uint8_t* weight, int32_t weight_byte_offset, + * int32_t input_dim, int32_t output_dim, + * float* output, int32_t output_offset); + * + * The C kernel decodes the ggml-canonical Q5_1 block (FP16 scale + FP16 min + * + 4-byte `qh` high-bit plane + 16 packed bytes, split nibble layout) with + * `d * code + m` affine dequant; on AArch64 a plain-NEON body (no dotprod + * requirement) expands the high-bit plane with `vtstq_u8` and FMAs the + * unsigned codes against the activation lanes. + * + * Numerical parity vs [ScalarQ5_1MatmulKernel] is asserted by + * `NativeQ5_1MatmulKernelParityTest` within the same `1e-2 * + * blocksPerInputDim` band the Panama parity uses. + */ +internal object NativeQ5_1MatmulKernel : Q5_1MatmulKernel { + + fun isAvailable(): Boolean = handle != null + + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) { + require(inputDim % BLOCK_SIZE == 0) { + "NativeQ5_1MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" + } + if (outputDim == 0) return + + val mh = handle + ?: error("NativeQ5_1MatmulKernel.matmul invoked while native library unavailable") + + val blocksPerInputDim = inputDim / BLOCK_SIZE + val inputReachFloats = if (inputDim == 0) 0 else inputOffset + inputDim + val weightReachBytes = if (inputDim == 0 || outputDim == 0) 0 + else weightByteOffset + blocksPerInputDim * outputDim * BYTES_PER_BLOCK + val outputReachFloats = outputOffset + outputDim + + Arena.ofConfined().use { arena -> + val fAlign = ValueLayout.JAVA_FLOAT.byteAlignment() + val bAlign = ValueLayout.JAVA_BYTE.byteAlignment() + + val inputSeg: MemorySegment = if (inputReachFloats > 0) + arena.allocate(inputReachFloats.toLong() * java.lang.Float.BYTES, fAlign) + else MemorySegment.NULL + val weightSeg: MemorySegment = if (weightReachBytes > 0) + arena.allocate(weightReachBytes.toLong(), bAlign) + else MemorySegment.NULL + val outputSeg: MemorySegment = + arena.allocate(outputReachFloats.toLong() * java.lang.Float.BYTES, fAlign) + + if (inputReachFloats > 0) { + MemorySegment.copy(input, 0, inputSeg, ValueLayout.JAVA_FLOAT, 0L, inputReachFloats) + } + if (weightReachBytes > 0) { + MemorySegment.copy(weight, 0, weightSeg, ValueLayout.JAVA_BYTE, 0L, weightReachBytes) + } + + mh.invoke( + inputSeg, inputOffset, + weightSeg, weightByteOffset, + inputDim, outputDim, + outputSeg, outputOffset, + ) + + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, outputReachFloats) + } + } + + private const val BLOCK_SIZE = 32 + private const val BYTES_PER_BLOCK = 24 + + private val handle: MethodHandle? by lazy { + val lookup = NativeLibraryLoader.lookup() ?: return@lazy null + val symbol = lookup.find("skainet_q5_1_matmul").orElse(null) ?: return@lazy null + val descriptor = FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, // input + ValueLayout.JAVA_INT, // input_offset + ValueLayout.ADDRESS, // weight + ValueLayout.JAVA_INT, // weight_byte_offset + ValueLayout.JAVA_INT, // input_dim + ValueLayout.JAVA_INT, // output_dim + ValueLayout.ADDRESS, // output + ValueLayout.JAVA_INT, // output_offset + ) + runCatching { Linker.nativeLinker().downcallHandle(symbol, descriptor) }.getOrNull() + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt index 2326f2890..42f7ee033 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelSupportMatrixTest.kt @@ -40,9 +40,9 @@ class KernelSupportMatrixTest { Tier("panama-vector", 50, setOf("JVM", "Android"), setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q6_K", "Q5_K", "Q5_1", "Q5_0")), Tier("native-ffm", 100, setOf("JVM"), - setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q5_K")), + setOf("Float32", "BFloat16", "Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q5_1", "Q5_0")), Tier("native-jni", 100, setOf("Android"), - setOf("Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K")), + setOf("Q8_0", "Q4_0", "Q4_K", "Q5_K", "Q6_K", "Q5_1", "Q5_0")), ) private fun best(fmt: String, platform: String, tiers: List): String? = diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernelParityTest.kt new file mode 100644 index 000000000..9bd4e3c19 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernelParityTest.kt @@ -0,0 +1,118 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Numerical parity tests for [NativeQ5_0MatmulKernel] against + * [ScalarQ5_0MatmulKernel]. Same FP16 scale decode + split-layout + * `nibble | (fifth_bit << 4)` code assembly in both kernels; the C body + * folds the `- 16` re-centring algebraically (`d * (dot - 16 * sum(x))`), + * so differences come from FMA + reordered-reduction only. + * + * Tolerance: `1e-2 * blocksPerInputDim` (matches the Panama / Q4_0 + * parity convention). + */ +class NativeQ5_0MatmulKernelParityTest { + + private val blockSize = 32 + private val bytesPerBlock = 22 + + @BeforeTest + fun checkAvailable() { + assertTrue( + NativeQ5_0MatmulKernel.isAvailable(), + "Native Q5_0 kernel must be available — bundled libskainet_kernels missing or " + + "skainet_q5_0_matmul symbol unresolved", + ) + } + + private fun randomQ5_0Bytes(blocksPerInputDim: Int, outputDim: Int, seed: Int): ByteArray { + val rng = Random(seed) + val numBlocks = blocksPerInputDim * outputDim + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + val base = block * bytesPerBlock + bytes[base + 0] = 0x00.toByte() + bytes[base + 1] = 0x22.toByte() // FP16 ~ 7.6e-3, comfortably finite + non-zero + } + return bytes + } + + private fun assertParity( + inputDim: Int, + outputDim: Int, + seed: Int, + tolPerBlock: Float = 1e-2f, + ) { + val blocksPerInputDim = inputDim / blockSize + val rng = Random(seed) + val input = FloatArray(inputDim) { rng.nextFloat() - 0.5f } + val weight = randomQ5_0Bytes(blocksPerInputDim, outputDim, seed) + val outScalar = FloatArray(outputDim) + val outNative = FloatArray(outputDim) + + ScalarQ5_0MatmulKernel.matmul(input, 0, weight, 0, inputDim, outputDim, outScalar, 0) + NativeQ5_0MatmulKernel.matmul(input, 0, weight, 0, inputDim, outputDim, outNative, 0) + + val tol = (tolPerBlock * blocksPerInputDim.coerceAtLeast(1)).coerceAtLeast(tolPerBlock) + for (i in outScalar.indices) { + val diff = abs(outScalar[i] - outNative[i]) + assertTrue( + diff <= tol, + "mismatch at $i: scalar=${outScalar[i]} native=${outNative[i]} diff=$diff tol=$tol", + ) + } + } + + @Test fun single_block_single_output_matches_scalar() = + assertParity(inputDim = 32, outputDim = 1, seed = 1) + + @Test fun single_block_multiple_outputs_matches_scalar() = + assertParity(inputDim = 32, outputDim = 7, seed = 2) + + @Test fun multiple_blocks_single_output_matches_scalar() = + assertParity(inputDim = 256, outputDim = 1, seed = 3) + + @Test fun llm_typical_attention_proj_matches_scalar() = + assertParity(inputDim = 512, outputDim = 512, seed = 4) + + @Test fun llm_typical_ffn_proj_matches_scalar() = + assertParity(inputDim = 256, outputDim = 1024, seed = 5) + + @Test fun rejects_non_block_aligned_input_dim() { + assertFailsWith { + NativeQ5_0MatmulKernel.matmul( + FloatArray(31), 0, + ByteArray(bytesPerBlock), 0, + 31, 1, + FloatArray(1), 0, + ) + } + } + + @Test fun zero_input_dim_zeros_output() { + val out = FloatArray(5) { 9f } + NativeQ5_0MatmulKernel.matmul( + FloatArray(0), 0, + ByteArray(0), 0, + 0, 5, + out, 0, + ) + for (v in out) assertEquals(0f, v, "output should be zeroed for inputDim=0") + } + + @Test fun provider_returns_native_q5_0_when_available() { + val kernel = NativeKernelProvider.matmulQ5_0() + assertTrue( + kernel === NativeQ5_0MatmulKernel, + "Provider must hand out the native Q5_0 kernel when bundled lib is loaded", + ) + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernelParityTest.kt new file mode 100644 index 000000000..4537a7257 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernelParityTest.kt @@ -0,0 +1,120 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Numerical parity tests for [NativeQ5_1MatmulKernel] against + * [ScalarQ5_1MatmulKernel]. Same FP16 `d`/`m` decode + split-layout + * `nibble | (fifth_bit << 4)` code assembly in both kernels; the C body + * folds the affine dequant algebraically (`d * dot + m * sum(x)`), so + * differences come from FMA + reordered-reduction only. + * + * Tolerance: `1e-2 * blocksPerInputDim` (matches the Panama / Q4_0 + * parity convention). + */ +class NativeQ5_1MatmulKernelParityTest { + + private val blockSize = 32 + private val bytesPerBlock = 24 + + @BeforeTest + fun checkAvailable() { + assertTrue( + NativeQ5_1MatmulKernel.isAvailable(), + "Native Q5_1 kernel must be available — bundled libskainet_kernels missing or " + + "skainet_q5_1_matmul symbol unresolved", + ) + } + + private fun randomQ5_1Bytes(blocksPerInputDim: Int, outputDim: Int, seed: Int): ByteArray { + val rng = Random(seed) + val numBlocks = blocksPerInputDim * outputDim + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + val base = block * bytesPerBlock + bytes[base + 0] = 0x00.toByte() + bytes[base + 1] = 0x22.toByte() // d: FP16 ~ 7.6e-3, finite + non-zero + bytes[base + 2] = 0x00.toByte() + bytes[base + 3] = 0x9E.toByte() // m: FP16 ~ -7.6e-3 (negative min exercised) + } + return bytes + } + + private fun assertParity( + inputDim: Int, + outputDim: Int, + seed: Int, + tolPerBlock: Float = 1e-2f, + ) { + val blocksPerInputDim = inputDim / blockSize + val rng = Random(seed) + val input = FloatArray(inputDim) { rng.nextFloat() - 0.5f } + val weight = randomQ5_1Bytes(blocksPerInputDim, outputDim, seed) + val outScalar = FloatArray(outputDim) + val outNative = FloatArray(outputDim) + + ScalarQ5_1MatmulKernel.matmul(input, 0, weight, 0, inputDim, outputDim, outScalar, 0) + NativeQ5_1MatmulKernel.matmul(input, 0, weight, 0, inputDim, outputDim, outNative, 0) + + val tol = (tolPerBlock * blocksPerInputDim.coerceAtLeast(1)).coerceAtLeast(tolPerBlock) + for (i in outScalar.indices) { + val diff = abs(outScalar[i] - outNative[i]) + assertTrue( + diff <= tol, + "mismatch at $i: scalar=${outScalar[i]} native=${outNative[i]} diff=$diff tol=$tol", + ) + } + } + + @Test fun single_block_single_output_matches_scalar() = + assertParity(inputDim = 32, outputDim = 1, seed = 1) + + @Test fun single_block_multiple_outputs_matches_scalar() = + assertParity(inputDim = 32, outputDim = 7, seed = 2) + + @Test fun multiple_blocks_single_output_matches_scalar() = + assertParity(inputDim = 256, outputDim = 1, seed = 3) + + @Test fun llm_typical_attention_proj_matches_scalar() = + assertParity(inputDim = 512, outputDim = 512, seed = 4) + + @Test fun llm_typical_ffn_proj_matches_scalar() = + assertParity(inputDim = 256, outputDim = 1024, seed = 5) + + @Test fun rejects_non_block_aligned_input_dim() { + assertFailsWith { + NativeQ5_1MatmulKernel.matmul( + FloatArray(31), 0, + ByteArray(bytesPerBlock), 0, + 31, 1, + FloatArray(1), 0, + ) + } + } + + @Test fun zero_input_dim_zeros_output() { + val out = FloatArray(5) { 9f } + NativeQ5_1MatmulKernel.matmul( + FloatArray(0), 0, + ByteArray(0), 0, + 0, 5, + out, 0, + ) + for (v in out) assertEquals(0f, v, "output should be zeroed for inputDim=0") + } + + @Test fun provider_returns_native_q5_1_when_available() { + val kernel = NativeKernelProvider.matmulQ5_1() + assertTrue( + kernel === NativeQ5_1MatmulKernel, + "Provider must hand out the native Q5_1 kernel when bundled lib is loaded", + ) + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnKernelProvider.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnKernelProvider.kt index 5bc47defd..42ffdd4ed 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnKernelProvider.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnKernelProvider.kt @@ -10,10 +10,14 @@ import sk.ainet.backend.api.kernel.KernelRegistry import sk.ainet.backend.api.kernel.Q4KMatmulKernel import sk.ainet.backend.api.kernel.Q4_0MatmulKernel import sk.ainet.backend.api.kernel.Q5KMatmulKernel +import sk.ainet.backend.api.kernel.Q5_0MatmulKernel +import sk.ainet.backend.api.kernel.Q5_1MatmulKernel import sk.ainet.backend.api.kernel.Q6KMatmulKernel import sk.ainet.backend.api.kernel.Q8_0MatmulKernel import sk.ainet.kernels.cinterop.skainet_q4_0_matmul import sk.ainet.kernels.cinterop.skainet_q4k_matmul +import sk.ainet.kernels.cinterop.skainet_q5_0_matmul +import sk.ainet.kernels.cinterop.skainet_q5_1_matmul import sk.ainet.kernels.cinterop.skainet_q6k_matmul import sk.ainet.kernels.cinterop.skainet_q8_0_matmul @@ -25,8 +29,8 @@ import sk.ainet.kernels.cinterop.skainet_q8_0_matmul * * **Registration is manual on K/N** (no `ServiceLoader`): a consumer calls * [installNativeKernels] once at startup. [Q5KMatmulKernel] (the FunctionGemma - * Q5_K_M hot path) plus Q4_K / Q6_K / Q8_0 / Q4_0 are wired; the rest cascade to - * the scalar provider. + * Q5_K_M hot path) plus Q4_K / Q6_K / Q8_0 / Q4_0 / Q5_0 / Q5_1 are wired; the + * rest cascade to the scalar provider. */ @OptIn(ExperimentalForeignApi::class) public object NativeKnKernelProvider : KernelProvider { @@ -45,6 +49,8 @@ public object NativeKnKernelProvider : KernelProvider { override fun matmulQ6K(): Q6KMatmulKernel = NativeKnQ6KMatmulKernel override fun matmulQ8_0(): Q8_0MatmulKernel = NativeKnQ8_0MatmulKernel override fun matmulQ4_0(): Q4_0MatmulKernel = NativeKnQ4_0MatmulKernel + override fun matmulQ5_0(): Q5_0MatmulKernel = NativeKnQ5_0MatmulKernel + override fun matmulQ5_1(): Q5_1MatmulKernel = NativeKnQ5_1MatmulKernel } /** @@ -156,3 +162,51 @@ public object NativeKnQ4_0MatmulKernel : Q4_0MatmulKernel { } } } } } + +@OptIn(ExperimentalForeignApi::class) +public object NativeKnQ5_0MatmulKernel : Q5_0MatmulKernel { + private const val BLOCK_SIZE = 32 + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) { + require(inputDim % BLOCK_SIZE == 0) { + "NativeKnQ5_0MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" + } + if (outputDim == 0 || inputDim == 0) return + input.usePinned { i -> weight.usePinned { w -> output.usePinned { o -> + skainet_q5_0_matmul( + i.addressOf(0), inputOffset, + w.addressOf(0).reinterpret(), weightByteOffset, + inputDim, outputDim, + o.addressOf(0), outputOffset, + ) + } } } + } +} + +@OptIn(ExperimentalForeignApi::class) +public object NativeKnQ5_1MatmulKernel : Q5_1MatmulKernel { + private const val BLOCK_SIZE = 32 + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) { + require(inputDim % BLOCK_SIZE == 0) { + "NativeKnQ5_1MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" + } + if (outputDim == 0 || inputDim == 0) return + input.usePinned { i -> weight.usePinned { w -> output.usePinned { o -> + skainet_q5_1_matmul( + i.addressOf(0), inputOffset, + w.addressOf(0).reinterpret(), weightByteOffset, + inputDim, outputDim, + o.addressOf(0), outputOffset, + ) + } } } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnQ5_0MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnQ5_0MatmulKernelParityTest.kt new file mode 100644 index 000000000..e4336a301 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnQ5_0MatmulKernelParityTest.kt @@ -0,0 +1,73 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * Proves the Kotlin/Native cinterop path: [NativeKnQ5_0MatmulKernel] (calling + * the C `skainet_q5_0_matmul` via cinterop, linked from libskainet_kernels.a) + * must agree with the commonMain [ScalarQ5_0MatmulKernel] reference within + * FMA + `-ffast-math` reassociation tolerance. + * + * Runs on linuxX64 (host archive: scalar/auto-vectorized) AND linuxArm64 + * (cross-built archive: plain-NEON body — high-bit plane expanded with + * `vtstq_u8`, no dotprod requirement), so the aarch64 run checks the + * `SKAINET_HAVE_NEON` path in q5_0_matmul.c against the scalar reference. + * Q5_0 blocks are 32 elements / 22 bytes (FP16 `d` + 4-byte `qh` high-bit + * plane + 16 bytes of split-layout nibbles); full-range random `qh`/`qs` + * bytes exercise both nibble lanes, all 32 high-bit positions and the + * `- 16` re-centring. + */ +class NativeKnQ5_0MatmulKernelParityTest { + + private val blockSize = 32 + private val bytesPerBlock = 22 + + private fun randomQ5_0Bytes(numBlocks: Int, seed: Int): ByteArray { + val rng = Random(seed) + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + val base = block * bytesPerBlock + // 0x3C00 == 1.0f16 for the per-block scale so dequant stays finite. + bytes[base + 0] = 0x00.toByte() + bytes[base + 1] = 0x3C.toByte() + } + return bytes + } + + private fun assertParity(inputDim: Int, outputDim: Int, seed: Int, tol: Float) { + val numBlocks = (inputDim / blockSize) * outputDim + val packed = randomQ5_0Bytes(numBlocks, seed) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + + val refOut = FloatArray(outputDim) + ScalarQ5_0MatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, refOut, 0) + + val knOut = FloatArray(outputDim) + NativeKnQ5_0MatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, knOut, 0) + + for (o in 0 until outputDim) { + val diff = abs(refOut[o] - knOut[o]) + val rel = diff / (abs(refOut[o]) + 1e-9f) + assertTrue( + diff <= tol || rel < 1e-4f, + "row $o diverged: scalar=${refOut[o]} cinterop=${knOut[o]} diff=$diff rel=$rel tol=$tol", + ) + } + } + + @Test + fun single_block_single_row() = assertParity(32, 1, 42, 1e-2f) + + @Test + fun single_block_multi_row() = assertParity(32, 16, 7, 1e-2f) + + @Test + fun multi_block_multi_row() = assertParity(1024, 64, 123, 2e-1f) + + @Test + fun llm_typical_shape() = assertParity(4096, 64, 999, 2e0f) +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnQ5_1MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnQ5_1MatmulKernelParityTest.kt new file mode 100644 index 000000000..fc5fe671c --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnQ5_1MatmulKernelParityTest.kt @@ -0,0 +1,76 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * Proves the Kotlin/Native cinterop path: [NativeKnQ5_1MatmulKernel] (calling + * the C `skainet_q5_1_matmul` via cinterop, linked from libskainet_kernels.a) + * must agree with the commonMain [ScalarQ5_1MatmulKernel] reference within + * FMA + `-ffast-math` reassociation tolerance. + * + * Runs on linuxX64 (host archive: scalar/auto-vectorized) AND linuxArm64 + * (cross-built archive: plain-NEON body — high-bit plane expanded with + * `vtstq_u8`, no dotprod requirement), so the aarch64 run checks the + * `SKAINET_HAVE_NEON` path in q5_1_matmul.c against the scalar reference. + * Q5_1 blocks are 32 elements / 24 bytes (FP16 `d` + FP16 `m` + 4-byte `qh` + * high-bit plane + 16 bytes of split-layout nibbles); `d` and `m` are both + * pinned to 1.0f16 so the affine `d*code + m` dequant stays finite while + * full-range random `qh`/`qs` bytes exercise both nibble lanes and all 32 + * high-bit positions. + */ +class NativeKnQ5_1MatmulKernelParityTest { + + private val blockSize = 32 + private val bytesPerBlock = 24 + + private fun randomQ5_1Bytes(numBlocks: Int, seed: Int): ByteArray { + val rng = Random(seed) + val bytes = ByteArray(numBlocks * bytesPerBlock) + rng.nextBytes(bytes) + for (block in 0 until numBlocks) { + val base = block * bytesPerBlock + // 0x3C00 == 1.0f16 for the per-block scale AND min so dequant stays finite. + bytes[base + 0] = 0x00.toByte() + bytes[base + 1] = 0x3C.toByte() + bytes[base + 2] = 0x00.toByte() + bytes[base + 3] = 0x3C.toByte() + } + return bytes + } + + private fun assertParity(inputDim: Int, outputDim: Int, seed: Int, tol: Float) { + val numBlocks = (inputDim / blockSize) * outputDim + val packed = randomQ5_1Bytes(numBlocks, seed) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + + val refOut = FloatArray(outputDim) + ScalarQ5_1MatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, refOut, 0) + + val knOut = FloatArray(outputDim) + NativeKnQ5_1MatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, knOut, 0) + + for (o in 0 until outputDim) { + val diff = abs(refOut[o] - knOut[o]) + val rel = diff / (abs(refOut[o]) + 1e-9f) + assertTrue( + diff <= tol || rel < 1e-4f, + "row $o diverged: scalar=${refOut[o]} cinterop=${knOut[o]} diff=$diff rel=$rel tol=$tol", + ) + } + } + + @Test + fun single_block_single_row() = assertParity(32, 1, 42, 1e-2f) + + @Test + fun single_block_multi_row() = assertParity(32, 16, 7, 1e-2f) + + @Test + fun multi_block_multi_row() = assertParity(1024, 64, 123, 2e-1f) + + @Test + fun llm_typical_shape() = assertParity(4096, 64, 999, 2e0f) +}