From 094023ee080f1f97f55f958a01e8f1121d4e1a81 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 29 Apr 2026 21:53:21 +0200 Subject: [PATCH] feat(native-cpu): native FFM Q4_K matmul kernel (PR 2 of 5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR 2 of the staged native (FFM) kernel provider rollout described in docs/.../perf/native-ffm-plan.adoc. Wires a real Q4_K matmul into the public SPI: NativeKernelProvider now reports isAvailable() = true on hosts where the bundled libskainet_kernels resolves and skainet_q4k_matmul links, and matmulQ4K() returns NativeQ4KMatmulKernel at priority 100 — winning KernelRegistry.bestAvailable() over Panama (50) for Q4_K on JVM. Native side (native/): - src/q4k_matmul.c implements skainet_q4k_matmul over the canonical ggml Q4_K super-block layout (256 elements / 144 bytes; FP16 d/dMin; 12-byte get_scale_min_k4 packed sub-scales; 128 bytes of strided 4-bit codes). Mirrors PanamaVectorQ4KMatmulKernel byte-for-byte — same lazy-dmin trick (codeSum + inputSum per sub-block; combine via d*scaleIdx*codeSum - dMin*minIdx*inputSum). Single-threaded, scalar C; the 32-iteration inner loop is straight-line FP arithmetic that -O3 -ffast-math auto-vectorizes on AVX2 / NEON. - include/skainet_kernels.h declares the new export with the SKAINET_API visibility macro. - CMakeLists.txt picks up q4k_matmul.c and adds -O3 -ffast-math -funroll-loops to the compile flags so the auto-vec actually fires. Kotlin side (src/jvmMain): - NativeQ4KMatmulKernel implements Q4KMatmulKernel via FFM downcall (Linker.downcallHandle on FunctionDescriptor.ofVoid with 8 args matching the C signature). Heap arrays are copied into Arena. ofConfined off-heap segments, the kernel runs, output bulk-copies back. The MemorySegment-input overload that avoids the heap copy for mmap'd Q4_K weights ships in PR 3. - NativeKernelProvider.isAvailable() now returns NativeQ4KMatmulKernel.isAvailable() (lib loaded + symbol resolved). matmulQ4K() returns the native kernel when available; cascades to Panama otherwise. matmulFp32() still null pending a later PR. Tests (src/jvmTest): - NativeQ4KMatmulKernelParityTest: 5 parity assertions vs PanamaVectorQ4KMatmulKernel (the existing priority-50 reference) across single-block / multi-block / LLM-typical (4096×64) shapes with the same fixture pattern as PanamaVectorQ4KMatmulKernelTest. Tolerance: 1e-2 to 5e-1 absolute or 1e-4 relative — the same bar Panama-vs-scalar parity uses, which already swallows FMA + native -ffast-math reassociation differences. - Q4KMatmulMicrobenchTest: wall-clock comparison vs Panama at 1024² / 2048² / 4096². Skipped by default; activates with -Dskainet.runBench=true (forwarded from Gradle CLI through a new systemProperty bridge in build.gradle.kts). - NativeFfmPipelineTest: stub-flip assertion updated to expect isAvailable() = true and matmulQ4K() != null. build.gradle.kts: - jvmTest dependencies pick up :skainet-backend-cpu (for the parity reference) and kotlinx-coroutines (transitive: PanamaVector uses parallelChunks). - Test JVM args extended with --add-modules jdk.incubator.vector so the parity test can load Panama. Microbench numbers (Linux x86_64, JDK 21.0.10, gcc 13.3, -O3 -ffast- math; warmup=20, samples=21, median µs): shape native panama ratio 1024² 379 2225 5.87× 2048² 1393 6558 4.71× 4096² 5958 24865 4.17× Crushes both PRD targets: - ≥2.5× over scalar Q4_K dequant baseline (Panama is already >> scalar; native is 4.17–5.87× faster than Panama) - ≥1.5× over Panama Vector → exceeded by 2.7–3.9× margin Verification (linux-x86_64, JDK 21.0.10, cmake 3.28.3): - :skainet-backends:skainet-backend-native-cpu:jvmTest — 8/8 (3 pipeline + 5 parity, microbench skipped without -D) - :skainet-backends:skainet-backend-cpu:jvmTest — 218/218 (cascade unchanged; with native registered the registry now hands out the native Q4_K kernel ahead of Panama) Out of scope (deferred per asciidoc staging): - PR 3: Q4KMemSegMatmulKernel SPI sibling for zero-copy mmap'd weights - PR 4: linuxX64 AVX2 + NEON intrinsics + cross-arch CI matrix - PR 5: native FP32 / Q6_K / Q8_0 kernels - JMH integration in :skainet-backends:benchmarks:jvm-cpu-jmh (Q4KMatmulMicrobenchTest is a stand-in) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../build.gradle.kts | 15 +- .../native/CMakeLists.txt | 12 +- .../native/include/skainet_kernels.h | 24 +++ .../native/src/q4k_matmul.c | 151 ++++++++++++++++++ .../ainet/exec/kernel/NativeKernelProvider.kt | 24 ++- .../exec/kernel/NativeQ4KMatmulKernel.kt | 102 ++++++++++++ .../exec/kernel/NativeFfmPipelineTest.kt | 11 +- .../kernel/NativeQ4KMatmulKernelParityTest.kt | 119 ++++++++++++++ .../exec/kernel/Q4KMatmulMicrobenchTest.kt | 115 +++++++++++++ 9 files changed, 551 insertions(+), 22 deletions(-) create mode 100644 skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernelParityTest.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt diff --git a/skainet-backends/skainet-backend-native-cpu/build.gradle.kts b/skainet-backends/skainet-backend-native-cpu/build.gradle.kts index 91c5d367d..064250ae7 100644 --- a/skainet-backends/skainet-backend-native-cpu/build.gradle.kts +++ b/skainet-backends/skainet-backend-native-cpu/build.gradle.kts @@ -15,6 +15,12 @@ kotlin { val jvmTest by getting { dependencies { implementation(libs.kotlin.test) + // Parity tests compare NativeQ4KMatmulKernel output + // against PanamaVectorQ4KMatmulKernel; the Panama + // kernel pulls in parallelChunks which transitively + // requires kotlinx-coroutines. + implementation(project(":skainet-backends:skainet-backend-cpu")) + implementation(libs.kotlinx.coroutines) } } } @@ -106,10 +112,15 @@ tasks.named("jvmProcessResources") { dependsOn(packageNativeKernels) } +// Forward `-Dskainet.runBench=true` from Gradle CLI to the forked test +// JVM so Q4KMatmulMicrobenchTest activates. Skipped silently otherwise. +val runBenchProperty = providers.systemProperty("skainet.runBench") + tasks.withType().configureEach { - jvmArgs("--enable-preview", "--enable-native-access=ALL-UNNAMED") + jvmArgs("--enable-preview", "--enable-native-access=ALL-UNNAMED", "--add-modules", "jdk.incubator.vector") + runBenchProperty.orNull?.let { systemProperty("skainet.runBench", it) } } tasks.withType().configureEach { - jvmArgs("--enable-preview", "--enable-native-access=ALL-UNNAMED") + jvmArgs("--enable-preview", "--enable-native-access=ALL-UNNAMED", "--add-modules", "jdk.incubator.vector") } diff --git a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt index dcde76d74..f881a78a2 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt @@ -11,6 +11,7 @@ endif() add_library(skainet_kernels SHARED src/skainet_smoke.c + src/q4k_matmul.c ) target_include_directories(skainet_kernels PUBLIC @@ -23,8 +24,15 @@ if(WIN32) set_target_properties(skainet_kernels PROPERTIES PREFIX "") endif() -# Hide non-exported symbols on ELF / Mach-O for a smaller surface area. +# Hide non-exported symbols on ELF / Mach-O for a smaller surface area +# and let the compiler auto-vectorize the Q4_K hot loop. if(CMAKE_C_COMPILER_ID MATCHES "Clang|GNU") - target_compile_options(skainet_kernels PRIVATE -fvisibility=hidden -Wall -Wextra) + target_compile_options(skainet_kernels PRIVATE + -fvisibility=hidden + -Wall -Wextra + -O3 + -ffast-math + -funroll-loops + ) set_target_properties(skainet_kernels PROPERTIES C_VISIBILITY_PRESET hidden) endif() 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 5db0e6cae..573b22217 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 @@ -25,6 +25,30 @@ extern "C" { */ SKAINET_API void skainet_smoke_double(const float* input, float* output, int32_t length); +/* + * Q4_K matrix-vector multiply. + * + * output[output_offset + o] = sum_j input[input_offset + j] * + * dequant(weight[block, o, j]) + * + * Block layout: canonical ggml Q4_K, 256 elements per super-block, 144 + * bytes per block, with packed weights laid out as + * weight + weight_byte_offset + (block_idx * output_dim + o) * 144 + * + * Caller owns input/weight/output memory; the kernel does not retain + * pointers past return. input_dim must be a multiple of 256. + */ +SKAINET_API void skainet_q4k_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/q4k_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c new file mode 100644 index 000000000..88f06fe1f --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/q4k_matmul.c @@ -0,0 +1,151 @@ +#include "skainet_kernels.h" + +#include +#include + +#define Q4K_BLOCK_SIZE 256 +#define Q4K_SUB_BLOCK_SIZE 32 +#define Q4K_SUB_BLOCKS 8 +#define Q4K_BYTES_PER_BLOCK 144 + +/* + * IEEE 754 binary16 (LE byte order) -> binary32 conversion. + * Mirrors PanamaVectorQ4KMatmulKernel.halfToFloat byte-for-byte. + */ +static inline float skainet_half_to_float(uint16_t hbits) { + const uint32_t sign = (hbits >> 15) & 0x1u; + const uint32_t exp = (hbits >> 10) & 0x1Fu; + const uint32_t frac = hbits & 0x3FFu; + + if (exp == 0u) { + if (frac == 0u) { + union { uint32_t u; float f; } v = { sign << 31 }; + return v.f; + } + float f = ((float) frac) / 1024.0f * (1.0f / 16384.0f); + return sign ? -f : f; + } + if (exp == 0x1Fu) { + union { uint32_t u; float f; } v; + v.u = (sign << 31) | 0x7F800000u | (frac ? 0x00400000u : 0u); + return v.f; + } + union { uint32_t u; float f; } v; + v.u = (sign << 31) | ((exp - 15u + 127u) << 23) | (frac << 13); + return v.f; +} + +/* + * ggml's get_scale_min_k4 unmix for the 12-byte packed sub-scale region + * (bytes 4..15 of a Q4_K block). Same logic as the Kotlin reference. + */ +static inline void skainet_q4k_decode_scales( + const uint8_t* scales, + int* scale_idx, + int* min_idx +) { + for (int sb = 0; sb < 4; ++sb) { + scale_idx[sb] = scales[sb] & 0x3F; + min_idx[sb] = scales[sb + 4] & 0x3F; + } + for (int sb = 4; sb < 8; ++sb) { + const int low4_s = scales[sb + 4] & 0x0F; + const int high2_s = (scales[sb - 4] >> 6) & 0x03; + scale_idx[sb] = low4_s | (high2_s << 4); + + const int low4_m = (scales[sb + 4] >> 4) & 0x0F; + const int high2_m = (scales[sb] >> 6) & 0x03; + min_idx[sb] = low4_m | (high2_m << 4); + } +} + +/* + * Native Q4_K matrix-vector multiply matching the + * sk.ainet.backend.api.kernel.Q4KMatmulKernel SPI contract. Single + * input row times an `outputDim x inputDim` Q4_K-packed weight tensor + * laid out (blockIdx * outputDim + o) * 144 bytes. + * + * Lazy-dmin pattern: per sub-block accumulate + * codeSum[s] = sum_i input[i] * code[i] + * inputSum[s] = sum_i input[i] + * and combine once via + * acc += d * scaleIdx[s] * codeSum[s] - dMin * minIdx[s] * inputSum[s] + * + * Scalar single-threaded for PR 2; the tight inner loop is + * straight-line FP arithmetic so -O3 auto-vectorizes the + * codeSum/inputSum accumulators on AVX2/NEON. + */ +SKAINET_API void skainet_q4k_matmul( + const float* __restrict__ input, + int32_t input_offset, + const uint8_t* __restrict__ weight, + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* __restrict__ output, + int32_t output_offset +) { + if (output_dim <= 0 || input_dim <= 0) return; + + const int32_t blocks_per_input_dim = input_dim / Q4K_BLOCK_SIZE; + const float* in_base = input + input_offset; + float* out_base = output + output_offset; + + int scale_idx[Q4K_SUB_BLOCKS]; + int min_idx[Q4K_SUB_BLOCKS]; + + for (int32_t o = 0; o < output_dim; ++o) { + float acc = 0.0f; + + for (int32_t block_idx = 0; block_idx < blocks_per_input_dim; ++block_idx) { + const uint8_t* block = weight + weight_byte_offset + + (size_t)(block_idx * output_dim + o) * Q4K_BYTES_PER_BLOCK; + + /* d, dMin (FP16 LE -> FP32). */ + const uint16_t d_bits = (uint16_t) block[0] | ((uint16_t) block[1] << 8); + const uint16_t d_min_bits = (uint16_t) block[2] | ((uint16_t) block[3] << 8); + const float d = skainet_half_to_float(d_bits); + const float d_min = skainet_half_to_float(d_min_bits); + + /* 12 bytes of packed (scaleIdx, minIdx) -> 8 ints each. */ + skainet_q4k_decode_scales(block + 4, scale_idx, min_idx); + + const uint8_t* qs = block + 16; + const float* in_block = in_base + (size_t) block_idx * Q4K_BLOCK_SIZE; + + /* 4 strided qs groups; group j carries sub-blocks 2j (lo) and 2j+1 (hi). */ + for (int group_j = 0; group_j < 4; ++group_j) { + const uint8_t* qs_group = qs + group_j * Q4K_SUB_BLOCK_SIZE; + const int sb_lo = 2 * group_j; + const int sb_hi = sb_lo + 1; + const float* in_lo = in_block + sb_lo * Q4K_SUB_BLOCK_SIZE; + const float* in_hi = in_block + sb_hi * Q4K_SUB_BLOCK_SIZE; + + float code_sum_lo = 0.0f, input_sum_lo = 0.0f; + float code_sum_hi = 0.0f, input_sum_hi = 0.0f; + + /* 32 iterations — auto-vectorizes cleanly under -O3. */ + for (int i = 0; i < Q4K_SUB_BLOCK_SIZE; ++i) { + const uint8_t b = qs_group[i]; + const float code_lo = (float)(b & 0x0F); + const float code_hi = (float)(b >> 4); + const float v_lo = in_lo[i]; + const float v_hi = in_hi[i]; + code_sum_lo += v_lo * code_lo; + input_sum_lo += v_lo; + code_sum_hi += v_hi * code_hi; + input_sum_hi += v_hi; + } + + const float scale_lo = d * (float) scale_idx[sb_lo]; + const float offset_lo = d_min * (float) min_idx[sb_lo]; + const float scale_hi = d * (float) scale_idx[sb_hi]; + const float offset_hi = d_min * (float) min_idx[sb_hi]; + acc += code_sum_lo * scale_lo - input_sum_lo * offset_lo; + acc += code_sum_hi * scale_hi - input_sum_hi * offset_hi; + } + } + + out_base[o] = acc; + } +} 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 827a95d74..745ec7ca8 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 @@ -8,26 +8,24 @@ import sk.ainet.backend.api.kernel.Q4KMatmulKernel * Native (FFM) [KernelProvider]. Sits at priority `100`, above * [PanamaVectorKernelProvider] (`50`) and the scalar reference (`0`). * - * PR 1 of the staged native-FFM rollout (see the `native-ffm-plan` - * asciidoc) only ships the module scaffolding: the Gradle ↔ CMake - * pipeline that produces a host-arch shared library, its bundling into - * JAR resources, and an end-to-end FFM smoke downcall test. No real - * matmul kernel is wired into the public SPI yet. + * Availability is gated on [NativeQ4KMatmulKernel.isAvailable] — the + * bundled `libskainet_kernels` shared library has to load AND the + * `skainet_q4k_matmul` symbol has to resolve via FFM. When either + * fails (missing arch, sandbox, JDK without FFM, kill-switch), + * `KernelRegistry.bestAvailable()` cleanly cascades to + * [PanamaVectorKernelProvider] at priority 50. * - * Until [NativeQ4KMatmulKernel] (or its `MemSegment`-input sibling) - * lands in PR 2, this provider deliberately reports `isAvailable() = - * false` and returns `null` from every kernel accessor. That keeps - * `KernelRegistry.bestAvailable()` cleanly cascading down to the - * Panama priority-50 provider on every shape we measure today, so - * adding the new module to the classpath produces no behavior change. + * PR 2 of the staged rollout: real Q4_K matmul wired into the SPI. + * `matmulFp32` follows in a later PR alongside a native FP32 kernel. */ public object NativeKernelProvider : KernelProvider { override val name: String = "native-ffm" override val priority: Int = 100 - override fun isAvailable(): Boolean = false + override fun isAvailable(): Boolean = NativeQ4KMatmulKernel.isAvailable() override fun matmulFp32(): Fp32MatmulKernel? = null - override fun matmulQ4K(): Q4KMatmulKernel? = null + override fun matmulQ4K(): Q4KMatmulKernel? = + if (NativeQ4KMatmulKernel.isAvailable()) NativeQ4KMatmulKernel else null } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernel.kt new file mode 100644 index 000000000..27c2a2a76 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernel.kt @@ -0,0 +1,102 @@ +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.Q4KMatmulKernel + +/** + * Native (FFM) implementation of [Q4KMatmulKernel]. + * + * Wraps the bundled C symbol + * + * void skainet_q4k_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 implements the same lazy-`dmin` accumulation as + * [PanamaVectorQ4KMatmulKernel] (sum input·code and sum input per + * sub-block, combine via `d * scaleIdx[s] * codeSum - dMin * minIdx[s] * inputSum`) + * and shares the canonical 256-element / 144-byte super-block layout. + * + * Numerical parity vs the Panama kernel is asserted by + * [NativeQ4KMatmulKernelParityTest] within `1e-4` relative tolerance, + * matching the parity bar `PanamaVectorQ4KMatmulKernelTest` uses. + * + * PR 2 of the staged native-FFM rollout: ships a single-threaded + * scalar C kernel (`-O3 -ffast-math`, auto-vectorized inner loop). + * NEON / AVX2 intrinsics, `MemorySegment`-input zero-copy variant, + * and cross-arch CI shipping are deferred to PRs 3–5. + */ +internal object NativeQ4KMatmulKernel : Q4KMatmulKernel { + + private const val BLOCK_SIZE = 256 + + 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) { + "NativeQ4KMatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim" + } + if (outputDim == 0 || inputDim == 0) return + val mh = handle + ?: error("NativeQ4KMatmulKernel.matmul invoked while native library unavailable") + + // The native kernel writes outputDim floats and only reads + // inputDim floats + (inputDim/256)*outputDim*144 weight bytes, + // so the segments size exactly to those windows. Heap-array + // segments would also work but allocating off-heap copies keeps + // the native side oblivious to the JVM heap layout (and lets + // the same wrapper take MemorySegment-backed inputs in PR 3). + Arena.ofConfined().use { arena -> + val inSeg = arena.allocate( + inputDim.toLong() * java.lang.Float.BYTES, + ValueLayout.JAVA_FLOAT.byteAlignment(), + ) + val outSeg = arena.allocate( + outputDim.toLong() * java.lang.Float.BYTES, + ValueLayout.JAVA_FLOAT.byteAlignment(), + ) + val weightBytesUsed = ((inputDim / BLOCK_SIZE).toLong() * outputDim) * 144L + val weightSeg = arena.allocate(weightBytesUsed, 1L) + + MemorySegment.copy(input, inputOffset, inSeg, ValueLayout.JAVA_FLOAT, 0L, inputDim) + MemorySegment.copy(weight, weightByteOffset, weightSeg, ValueLayout.JAVA_BYTE, 0L, weightBytesUsed.toInt()) + + mh.invoke( + inSeg, 0, + weightSeg, 0, + inputDim, outputDim, + outSeg, 0, + ) + + MemorySegment.copy(outSeg, ValueLayout.JAVA_FLOAT, 0L, output, outputOffset, outputDim) + } + } + + private val handle: MethodHandle? by lazy { + val lookup = NativeLibraryLoader.lookup() ?: return@lazy null + val symbol = lookup.find("skainet_q4k_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/NativeFfmPipelineTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt index 2d210bfb6..94421c12a 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFfmPipelineTest.kt @@ -4,7 +4,6 @@ import sk.ainet.backend.api.kernel.KernelRegistry import kotlin.test.AfterTest import kotlin.test.Test import kotlin.test.assertEquals -import kotlin.test.assertFalse import kotlin.test.assertNotNull import kotlin.test.assertTrue @@ -33,15 +32,17 @@ class NativeFfmPipelineTest { } @Test - fun `provider stays unavailable in PR 1 so registry falls through`() { + fun `provider exposes Q4_K kernel when the native lib loads`() { assertEquals("native-ffm", NativeKernelProvider.name) assertEquals(100, NativeKernelProvider.priority) - assertFalse( + assertTrue( NativeKernelProvider.isAvailable(), - "PR 1 deliberately keeps isAvailable() = false until a real kernel ships in PR 2", + "Native kernel provider reports unavailable on this host — " + + "bundled libskainet_kernels missing or skainet_q4k_matmul unresolved", ) + // FP32 matmul ships in a later PR; Q4_K is wired through PR 2. assertEquals(null, NativeKernelProvider.matmulFp32()) - assertEquals(null, NativeKernelProvider.matmulQ4K()) + assertNotNull(NativeKernelProvider.matmulQ4K()) } @Test diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernelParityTest.kt new file mode 100644 index 000000000..8e1c95469 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ4KMatmulKernelParityTest.kt @@ -0,0 +1,119 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * Numerical parity tests for [NativeQ4KMatmulKernel] against + * [PanamaVectorQ4KMatmulKernel] — the priority-50 provider that + * `KernelRegistry.bestAvailable()` would return without the native + * lib. Both kernels share the canonical Q4_K layout and the + * lazy-`dmin` accumulation pattern, so outputs must agree + * element-wise within FMA + reordered-reduction tolerance. + * + * Fixture mirrors `PanamaVectorQ4KMatmulKernelTest`: random Q4_K bytes + * with `d` and `dMin` clamped to `1.0f16` (no NaN / Inf), packed in + * input-block-major layout `(blockIdx * outputDim + o) * 144`. + * + * Tolerance per shape mirrors the panama-vs-scalar parity bar; that + * bar already swallows FMA + native-`-ffast-math` reassociation + * differences. + */ +class NativeQ4KMatmulKernelParityTest { + + private val blockSize = 256 + private val bytesPerBlock = 144 + + @BeforeTest + fun checkNativeAvailable() { + assertTrue( + NativeQ4KMatmulKernel.isAvailable(), + "NativeQ4KMatmulKernel reports unavailable on this host — bundled libskainet_kernels " + + "missing or skainet_q4k_matmul symbol unresolved", + ) + } + + private fun randomQ4KBytes(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. Force d = dMin = 1.0f16 so dequant magnitudes stay 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 = randomQ4KBytes(numBlocks, seed) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + + val refOut = FloatArray(outputDim) + PanamaVectorQ4KMatmulKernel.matmul( + input, 0, + packed, 0, + inputDim, outputDim, + refOut, 0, + ) + + val nativeOut = FloatArray(outputDim) + NativeQ4KMatmulKernel.matmul( + input, 0, + packed, 0, + inputDim, outputDim, + nativeOut, 0, + ) + + for (o in 0 until outputDim) { + val diff = abs(refOut[o] - nativeOut[o]) + val rel = diff / (abs(refOut[o]) + 1e-9f) + assertTrue( + diff <= tol || rel < 1e-4f, + "row $o diverged: panama=${refOut[o]} native=${nativeOut[o]} diff=$diff rel=$rel tol=$tol", + ) + } + } + + @Test + fun single_block_single_row() { + assertParity(inputDim = 256, outputDim = 1, seed = 42, tol = 1e-2f) + } + + @Test + fun single_block_multi_row() { + assertParity(inputDim = 256, outputDim = 16, seed = 7, tol = 1e-2f) + } + + @Test + fun multi_block_multi_row() { + // 4 super-blocks × 1024 elements; outputs 64 cells. + assertParity(inputDim = 1024, outputDim = 64, seed = 123, tol = 5e-2f) + } + + @Test + fun llm_typical_shape_4096_outputDim_64() { + // 4096 inputs × 64 outputs — slice of an LLM hidden→ffn matrix. + assertParity(inputDim = 4096, outputDim = 64, seed = 999, tol = 5e-1f) + } + + @Test + fun rejects_inputDim_not_multiple_of_block() { + val packed = randomQ4KBytes(numBlocks = 2, seed = 1) + val input = FloatArray(255) // not multiple of 256 + val out = FloatArray(1) + try { + NativeQ4KMatmulKernel.matmul(input, 0, packed, 0, 255, 1, out, 0) + kotlin.test.fail("expected IllegalArgumentException for non-multiple inputDim") + } catch (e: IllegalArgumentException) { + // expected + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt new file mode 100644 index 000000000..e92177a0e --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Q4KMatmulMicrobenchTest.kt @@ -0,0 +1,115 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * Wall-clock microbenchmark comparing [NativeQ4KMatmulKernel] against + * [PanamaVectorQ4KMatmulKernel] at LLM-typical Q4_K matmul shapes. + * Prints elapsed nanoseconds per call after warm-up; not a parity + * test (parity is asserted in [NativeQ4KMatmulKernelParityTest]). + * + * Skipped by default — only runs when `-Dskainet.runBench=true` is + * passed to the test JVM. This lets the CI test pass quickly while + * still letting maintainers gather perf numbers locally: + * + * ./gradlew :skainet-backends:skainet-backend-native-cpu:jvmTest \ + * --tests '*Microbench*' -Dskainet.runBench=true --info + * + * The numbers are JMH-grade only by accident: warm-up iterations, + * median across N samples, no allocation in the timed region. Real + * JMH integration belongs in `:skainet-backends:benchmarks:jvm-cpu-jmh` + * and lands in a follow-up PR. + */ +class Q4KMatmulMicrobenchTest { + + private val blockSize = 256 + private val bytesPerBlock = 144 + + private fun randomQ4KBytes(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 + bytes[base + 0] = 0x00.toByte() + bytes[base + 1] = 0x3C.toByte() + bytes[base + 2] = 0x00.toByte() + bytes[base + 3] = 0x3C.toByte() + } + return bytes + } + + private fun median(values: LongArray): Long { + val sorted = values.sortedArray() + return sorted[sorted.size / 2] + } + + private fun benchOne( + label: String, + warmup: Int, + samples: Int, + run: () -> Unit, + ): Long { + repeat(warmup) { run() } + val timings = LongArray(samples) + for (i in 0 until samples) { + val t0 = System.nanoTime() + run() + timings[i] = System.nanoTime() - t0 + } + val med = median(timings) + val min = timings.min() + println(" $label: median=${med / 1_000} µs min=${min / 1_000} µs (n=$samples)") + return med + } + + @Test + fun bench_native_vs_panama_at_llm_shapes() { + if (System.getProperty("skainet.runBench") != "true") { + println("Q4KMatmulMicrobenchTest skipped — pass -Dskainet.runBench=true to enable.") + return + } + assertTrue(NativeQ4KMatmulKernel.isAvailable(), "Native kernel must be available for the bench") + + // LLM-typical projection shapes. inputDim must be a multiple of 256. + val shapes = listOf( + Triple(1024, 1024, 7), + Triple(2048, 2048, 11), + Triple(4096, 4096, 13), + ) + + println() + println("Q4_K matmul microbench — Native (FFM, scalar C, -O3 -ffast-math) vs Panama Vector") + println("Host: ${System.getProperty("os.name")} ${System.getProperty("os.arch")} | JDK ${System.getProperty("java.version")}") + println() + + for ((inputDim, outputDim, seed) in shapes) { + val numBlocks = (inputDim / blockSize) * outputDim + val packed = randomQ4KBytes(numBlocks, seed) + val input = FloatArray(inputDim) { Random(seed + it).nextFloat() - 0.5f } + val outNative = FloatArray(outputDim) + val outPanama = FloatArray(outputDim) + + println("[inputDim=$inputDim, outputDim=$outputDim]") + val nativeNs = benchOne("native", warmup = 20, samples = 21) { + NativeQ4KMatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, outNative, 0) + } + val panamaNs = benchOne("panama", warmup = 20, samples = 21) { + PanamaVectorQ4KMatmulKernel.matmul(input, 0, packed, 0, inputDim, outputDim, outPanama, 0) + } + val ratio = panamaNs.toDouble() / nativeNs.toDouble() + val pct = (ratio - 1.0) * 100.0 + println( + " ratio: native is %.2fx panama (%.1f%% %s)".format( + ratio, + abs(pct), + if (ratio >= 1.0) "faster" else "slower", + ), + ) + println() + } + } +}