From d799ae53b63ccf6ab6457cf8ff6b60064f97f302 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 29 Apr 2026 07:33:34 +0200 Subject: [PATCH 1/2] docs: SIMD kernels, quantized SIMD, native FFM plan; arc42 architecture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three new explanation pages under docs/.../explanation/perf/ covering the M5 work that landed in 0.21.0: - simd-kernels.adoc — kernel SPI overview, FloatVector + FMA pattern, tile blocking, ServiceLoader auto-discovery + factory wrappers, KernelMatmulBench numbers (8.6×–10.8× over scalar at 256/512/1024 on Apple Silicon NEON). - quantized-simd-kernels.adoc — per-format pipelines (Q4_0, Q4_K, Q4_K MemSeg, Q6_K, Q8_0): ByteVector → AND/LSHR nibble extract → castShape(B2F) → fused FMA, plus the lazy-dmin trick and the per- format coverage matrix. - native-ffm-plan.adoc — recovers the FFM PRD content from git history (61962def:NATIVE_FFM_KERNEL_PROVIDER.md, dropped from 0.21.0 release per #566). Module layout, FFM binding pattern, staged delivery, success metrics, risks, trigger conditions. architecture.adoc grows from a 4-line stub to an arc42-style reference with focus on the 0.21.0 changes: - Building Block View — module table + kernel SPI ASCII diagram (commonMain api + jvmMain auto-discovery + jvmMain providers). - Runtime View — eager-execution flow from `ctx.ops.matmul` through `chooseQuantizedMatmul` / `chooseMatmul` to the SPI kernel, with the lazy provider resolution and fall-through pattern called out. - Architecture decisions table, quality requirements, risks (Vector API still incubator, no native provider yet, prior reverts). nav.adoc: register the three new explanation pages under .Explanation. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/modules/ROOT/nav.adoc | 3 + .../explanation/perf/native-ffm-plan.adoc | 278 +++++++++++++++++ .../perf/quantized-simd-kernels.adoc | 235 ++++++++++++++ .../pages/explanation/perf/simd-kernels.adoc | 244 +++++++++++++++ .../ROOT/pages/reference/architecture.adoc | 289 +++++++++++++++++- 5 files changed, 1043 insertions(+), 6 deletions(-) create mode 100644 docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc create mode 100644 docs/modules/ROOT/pages/explanation/perf/quantized-simd-kernels.adoc create mode 100644 docs/modules/ROOT/pages/explanation/perf/simd-kernels.adoc diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index a12d24b3e..5bbf20253 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -25,5 +25,8 @@ * xref:explanation/examples/index.adoc[Worked examples] ** xref:explanation/examples/matmul.adoc[Matrix multiplication examples] * xref:explanation/perf/jvm-cpu.adoc[JVM CPU performance] +* xref:explanation/perf/simd-kernels.adoc[How SIMD kernels are built] +* xref:explanation/perf/quantized-simd-kernels.adoc[How quantized SIMD kernels are built] +* xref:explanation/perf/native-ffm-plan.adoc[Plan: native FFM kernel provider] * xref:explanation/perf/java-25-cpu-backend.adoc[Java 25 CPU backend notes] * xref:explanation/issues/native-macos-accelerate-simd.adoc[Native macOS Accelerate SIMD issues] diff --git a/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc b/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc new file mode 100644 index 000000000..3e25e01b8 --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc @@ -0,0 +1,278 @@ += Plan: Native (FFM) Kernel Provider +:description: Where the JVM Vector kernels stop, what a native priority-100 provider would look like, and when to build it. + +This page is a *plan*, not shipped code. The intent is to capture +enough detail that the design doesn't drift between the time someone +decides to start the work and the moment a PR is opened. The earlier +content of this page lived briefly in `NATIVE_FFM_KERNEL_PROVIDER.md` +at the repo root and was removed on advice of "ship the release first, +keep the plan in docs"; this is its permanent home. + +== Where the JVM Vector kernels run out + +After the M5 milestone work landed (PRs #554–#565 across the 0.21.0 +release), every CPU matmul path goes through the kernel SPI — see +xref:explanation/perf/simd-kernels.adoc[] and +xref:explanation/perf/quantized-simd-kernels.adoc[]. The Panama Vector +provider runs at: + +* ~73 GFLOPS on FP32 4096² matmul (Apple Silicon NEON) +* ~73 GFLOPS on Q4_K 4096² matmul-vector (same regime; fused dequant +adds essentially zero cost on top of the FMA) + +That's already in the ggml NEON ballpark in absolute terms. But +ggml's hand-tuned NEON / AVX2 still outruns the JVM Vector API on: + +* dense FLOPs/cycle on shapes the Vector API can't tile-block +optimally (the 8×8×128 default is heuristic) +* AVX-512 VNNI fused INT8 dot products +* NEON `bf16` / `fp16` SDOT instructions +* future SVE / SME — none of which the Vector API exposes portably +today + +A native provider closes that gap and unlocks two follow-ons that +*can't* be built on the Vector API alone: + +. *M4 ↔ M5 zero-copy.* Mmap'd Q4_K weights stay as `MemorySegment` +views; a native kernel reads the same pages with no heap copy and +no staging buffer. +. *Hardware-specific lanes* unreachable from portable Vector code. + +== Provider shape + +[cols="1,1,1",options="header"] +|=== +| Priority | Provider | Status +| 0 | `ScalarKernelProvider` | shipped (PR #554) +| 50 | `PanamaVectorKernelProvider` | shipped (PRs #557, #560 + ServiceLoader #559) +| *100* | *`NativeKernelProvider` (FFM)* | *this plan* +|=== + +The `KernelRegistry.bestAvailable()` cascade means: when the native +lib loads, native wins; when it doesn't (sandbox, missing arch, JDK +without FFM, kill-switch flipped), Panama wins; on Native targets and +JS / Wasm where neither is available, scalar wins. No code change +above the registry layer. + +== Goals + +. *A `NativeKernelProvider` registered at priority 100* that on JDK +21+ wins `KernelRegistry.bestAvailable()` over Panama whenever the +native lib loads successfully. +. *A first concrete kernel: native Q4_K matmul.* It must: +.. take a `MemorySegment` for both input (FP32) and packed Q4_K +weights (canonical ggml layout — same as `Q4_KBlockTensorData` +and `matmulF32Q4_KMemSeg`); +.. produce numerically equivalent output to +`PanamaVectorQ4KMatmulKernel` within `1e-4` relative tolerance +(same parity bar `PanamaVectorQ4KMatmulKernelTest` uses); +.. clear *≥2.5×* over the prior Q4_K scalar dequant baseline — the +M5 success metric — on the bench shapes from +`QuantizedMatmulBench` (1024², 4096×1024, 4096²). +. *Optional follow-on kernels* — Q6_K, Q8_0, FP32 — share the build +system but each ship as a separate small PR. +. *One supported architecture for the first PR* (likely Apple +Silicon NEON since that's the development hardware in use), with a +clear extension path for `linuxX64` AVX2 / `linuxArm64` NEON. + +== Non-goals + +* *JNI.* The roadmap explicitly says "FFM not JNI". JNI's per-call +overhead and the global JNI lock are wrong for hot per-token +kernels; FFM (Java 22 stable, Java 21 preview) gives near-zero +overhead native calls and direct `MemorySegment` ABI. +* *Cross-compilation matrix on day one.* The first PR can ship just +one (host-arch) variant; CI cross-arch builds come later. +* *Replacing Panama.* Panama remains the priority-50 fallback for +environments that can't load native libs (sandboxes, Wasm, Native +targets, JDK without `jdk.incubator.vector`). +* *Distribution via pre-built native artifacts on Maven Central.* +Out of scope for the first PR — local build only. Publishing +classifier JARs comes in a separate plan. + +== Architecture + +=== Module layout + +[source] +---- +skainet-backends/ + skainet-backend-native-cpu/ # NEW + src/ + jvmMain/kotlin/sk/ainet/exec/kernel/ # Kotlin side + NativeKernelProvider.kt # priority=100, isAvailable()=libLoaded + NativeQ4KMatmulKernel.kt # implements Q4KMatmulKernel via FFM + NativeLibraryLoader.kt # System.loadLibrary, locate, version + jvmMain/resources/META-INF/services/ + sk.ainet.backend.api.kernel.KernelProvider # appends NativeKernelProviderFactory + jvmTest/kotlin/sk/ainet/exec/kernel/ + NativeQ4KMatmulKernelTest.kt # parity vs PanamaVectorQ4KMatmulKernel + native/ # native source tree + c/ + q4k_matmul.c # ggml-style hand-tuned kernel + q4k_matmul.h + CMakeLists.txt # or Bazel BUILD + build.gradle.kts # Gradle wrapper that invokes CMake +---- + +The native library compiles to a shared object (`libskainet_kernels.dylib` +on macOS, `.so` on Linux, `.dll` on Windows) and is packaged into the +module's resources for `System.loadLibrary` discovery. + +=== FFM binding pattern + +Single C entry point per kernel: + +[source,c] +---- +// q4k_matmul.h +void skainet_q4k_matmul( + const float* input, // FP32 input vector, length input_dim + const uint8_t* weight, // packed Q4_K bytes (canonical ggml layout) + int32_t weight_byte_offset, + int32_t input_dim, + int32_t output_dim, + float* output, // FP32 output, length output_dim + int32_t output_offset +); +---- + +Kotlin side: + +[source,kotlin] +---- +internal object NativeQ4KMatmulKernel : Q4KMatmulKernel { + private val handle: MethodHandle = run { + val arena = Arena.ofAuto() + val symbol = NativeLibraryLoader.lib.find("skainet_q4k_matmul").orElseThrow() + Linker.nativeLinker().downcallHandle( + symbol, + FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_INT, + ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, + ValueLayout.ADDRESS, ValueLayout.JAVA_INT, + ), + ) + } + + override fun matmul( + input: FloatArray, inputOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + output: FloatArray, outputOffset: Int, + ) { + // Heap arrays: pass via temporary off-heap MemorySegment + bulk copy, + // OR (preferred) overload with a MemorySegment-input variant for + // mmap'd weights to avoid the copy. + } +} +---- + +The cleaner path is to introduce a sibling `Q4KMemSegMatmulKernel` +SPI (mentioned as out-of-scope in PR #563) that takes `MemorySegment` +directly, and have the native provider implement *that* — no heap +copy. The `Q4KMatmulKernel` (`ByteArray`) variant can wrap the +MemSeg one with a temporary `Arena.ofConfined()` copy if needed for +legacy callers. + +=== Build system + +*Gradle + CMake* is the path of least resistance: + +* A new Gradle module (or hand-rolled `Exec` tasks) invokes CMake +for the native module's `build` task. +* Native artifacts land in `build/native//` and are copied +into `src/jvmMain/resources/native/-/` so +`System.loadLibrary` finds them. +* Kotlin compile depends on the native artifact being built first. + +The xnnpack backend already in the repo +(`skainet-backends/skainet-backend-xnnpack/`) demonstrates a similar +pattern — Gradle invokes CMake to build a native lib via cinterop. +*Reuse that template* rather than reinventing. + +== Staged delivery + +PRs in order, each independently mergeable: + +. *`skainet-backend-native-cpu` module scaffolding.* Gradle module, +`build.gradle.kts` wired to invoke CMake, a *trivial* C kernel +(e.g. just multiplies its first input by 2.0) to prove the FFM +pipeline end-to-end. `NativeKernelProvider` that's `isAvailable() += false` until the real kernel lands. Sets up CI artifact path on +host arch. +. *First real native kernel: Q4_K matmul (Apple Silicon NEON).* +Hand-tuned kernel, parity tests vs `PanamaVectorQ4KMatmulKernel`, +JMH bench variant added to `QuantizedMatmulBench`. +. *`Q4KMemSegMatmulKernel` SPI sibling + native variant.* Closes +the M4↔M5 zero-copy story for mmap'd weights. +. *`linuxX64` AVX2 variant + cross-arch CI build.* The +cross-compilation matrix story. +. *Optional: native FP32 matmul, native Q6_K, native Q8_0.* Same +shape as PRs 2–3, one per format. + +The first PR is the largest in scaffolding terms (~500–800 LoC of +build glue + 1 trivial kernel), but every subsequent PR is small and +template-able. + +== Success metrics + +* *PR 2 sign-off*: native Q4_K matmul on Apple Silicon clears *≥2.5×* +over the scalar Q4_K dequant-then-matmul baseline at 4096² (the M5 +milestone target). For reference: Panama Q4_K SIMD already exceeds +this metric (~73 GFLOPS, see +xref:explanation/perf/quantized-simd-kernels.adoc[]), so the bar is +"beats Panama by a meaningful margin", probably ≥1.5× over Panama. +* *PR 3 sign-off*: Q4_K MemSeg native path is faster than the Panama +Q4_K MemSeg path from PR #563, with no heap copy in the timed +region. +* *No regression on JVM-only environments* — when the native lib +fails to load (sandbox, missing arch, kill-switch), `bestAvailable()` +cleanly falls through to Panama, and existing tests / benches show +the same numbers as today. + +== Risks & open questions + +. *JDK 21 preview FFM vs JDK 22 stable.* FFM left preview in Java +22. The repo currently builds on JDK 21 with `--enable-preview +--add-modules jdk.incubator.vector`. Recommendation: stay on 21 +preview; flip to 22 in a separate toolchain-bump PR. +. *`MethodHandle` invocation overhead.* Even with FFM, each native +call has a small fixed cost (~µs). For the smallest matmul shapes +(e.g. 256² FP32) this could swamp the FLOPs win. Mitigation: route +small inputs to Panama and large inputs to native at the +registry/provider level, OR accept that the win is sized for +production-relevant shapes (4096²+). +. *Native code quality and maintenance.* Hand-tuned NEON / AVX2 in C +is harder to audit than Kotlin Vector API code. Mitigation: keep +kernels small (<300 LoC each), parity-test exhaustively, prefer +porting from ggml's reference (BSD-licensed, well-vetted) over +writing from scratch. +. *Distribution.* Native artifacts complicate Maven Central +publication (need `` per OS/arch). Not a blocker for +the first internal-use PR; a separate "publish native classifier +JARs" plan will be needed before community use. +. *Cross-arch CI cost.* Building NEON natively on Apple Silicon CI +plus AVX2 on linuxX64 plus Android NDK doubles or triples build +time. The xnnpack backend's existing CI matrix is a precedent — +reuse the same approach. +. *Native `MemorySegment` lifetime.* The Kotlin caller owns the +`Arena` for arrays it copies in. The native kernel must NOT retain +pointers past the FFM call return. Document this contract in +`NativeQ4KMatmulKernel.matmul` kdoc. + +== When to start + +Trigger conditions (any one): + +* Real workload demands the native ≥2.5× target (Panama Q4_K stops +being fast enough on a customer machine). +* A community contributor offers a hand-tuned NEON / AVX2 Q4_K +kernel that's measurably faster than Panama. +* A second M5 metric (e.g. SDPA throughput, training-loop +throughput) needs hand-tuned native code. + +Until then: *pause.* The Panama provider is doing the +milestone-equivalent work in absolute terms, and adding a native +build system is a meaningful complexity tax to take on +speculatively. diff --git a/docs/modules/ROOT/pages/explanation/perf/quantized-simd-kernels.adoc b/docs/modules/ROOT/pages/explanation/perf/quantized-simd-kernels.adoc new file mode 100644 index 000000000..9e625fffa --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/perf/quantized-simd-kernels.adoc @@ -0,0 +1,235 @@ += How Quantized SIMD Kernels Are Built +:description: ByteVector-based SIMD pipelines for Q4_0, Q4_K, Q6_K, Q8_0 quantized matmul on the JVM. + +For the broader kernel SPI story (`KernelProvider`, `KernelRegistry`, +ServiceLoader auto-discovery), see +xref:explanation/perf/simd-kernels.adoc[]. This page focuses on the +inner loops of the *quantized* matmul kernels — Q4_0, Q4_K, Q6_K, +Q8_0 — which is where most of the wall-clock time of an LLM decode +goes once the FP32 path is fast. + +== The general pipeline + +For a single `output[o] = Σ_j input[j] · dequant(weight[o, j])` cell, +the SIMD recipe is: + +. *Load codes.* `W` packed integers from the weight buffer as a +`ByteVector` — either via `ByteVector.fromArray` (heap `ByteArray`) +or `ByteVector.fromMemorySegment` (mmap'd weights, FFM). +. *Unpack.* For 4-bit codes: `byteVec.and(0x0F)` for low nibbles, +`byteVec.lanewise(LSHR, 4)` for high nibbles. Sign-correct as needed +(Q4_0 subtracts 8; Q4_K subtracts a per-sub-block min lazily; Q6_K +combines `ql + qh` then subtracts 32). +. *Widen + convert.* `castShape(floatSpecies, 0)` lane-widens the +byte vector to a `FloatVector` in one shape conversion (under the +hood: byte → int → float, but JIT'd as a single instruction +sequence on most targets). +. *Apply scale.* Multiply by the broadcast block scale (and +sub-scale, for K-quants). +. *Load input.* `W` floats from the FP32 input via +`FloatVector.fromArray`. +. *FMA.* `acc = inputVec.fma(weightFloatVec, acc)`. +. *Repeat across the block.* Reduce once per output cell at the end +via `acc.reduceLanes(ADD)`. + +The same skeleton drives every quantized kernel; the differences are +all in steps 1–4 (block layout, sign convention, scale recovery). + +== The four format pipelines + +=== Q8_0 — 32 elements / 34 bytes + +Single FP16 scale + 32 signed `int8` codes. The simplest case: codes +are already signed bytes, no nibble unpack. Pipeline: + +[source,kotlin] +---- +val byteVec = ByteVector.fromArray(byteSpeciesForFloat, codes, codesOffset + idx) +val codeVec = byteVec.castShape(floatSpecies, 0) as FloatVector // sign-extends +accVec = inputVec.mul(codeVec).add(accVec) +// final: (accVec.reduceLanes(ADD) + scalarTail) * scale +---- + +Q8_0 is the "gold reference" — the cleanest expression of the +pipeline. Q8_0 MemSeg (`matmulF32Q8_0MemSeg`) is the same loop with +`ByteVector.fromMemorySegment`. + +=== Q4_K — 256 elements / 144 bytes / 8 sub-blocks + +The format that matters most for current LLMs (Gemma 4 Q4_K_M, Llama, +Qwen). Block layout (canonical ggml): + +* bytes `[0, 2)`: `d` (super-block scale, FP16 LE) +* bytes `[2, 4)`: `dMin` (super-block min-scale, FP16 LE) +* bytes `[4, 16)`: 12 bytes packed `(scaleIdx, minIdx)` for 8 +sub-blocks via ggml's `get_scale_min_k4` mixing +* bytes `[16, 144)`: 128 bytes of 4-bit codes, *strided* in 4 groups +of 32 — each byte's lo nibble belongs to one sub-block, hi nibble to +the next sub-block over the same intra-group index. + +Per element: `dequant = code · scale[s] − offset[s]` where `scale[s] = +d · scaleIdx[s]` and `offset[s] = dMin · minIdx[s]`. + +==== The lazy-`dmin` trick + +A naive implementation subtracts `offset` from every element. Better: +linearity lets us track two running sums per sub-block: + +[source,text] +---- +codeSum[s] = Σ_i input[i] · code[i] (scaled later by scale[s]) +inputSum[s] = Σ_i input[i] (scaled later by offset[s]) +---- + +and combine once per sub-block as `acc += scale[s]·codeSum[s] − +offset[s]·inputSum[s]`. ggml's reference uses the same trick. + +==== The fused lo+hi load + +Because the canonical layout puts sub-block `2j` lo nibbles and +sub-block `2j+1` hi nibbles in the *same* 32-byte slab, a single +`ByteVector` load feeds both sub-block accumulators per chunk: + +[source,kotlin] +---- +val byteVec = ByteVector.fromArray(byteSpeciesForFloat, weight, qsRegion + idx) +val loBytes = byteVec.and(0x0F.toByte()) +val hiBytes = byteVec.lanewise(VectorOperators.LSHR, 4.toByte()) +val codeVecLo = loBytes.castShape(floatSpecies, 0) as FloatVector +val codeVecHi = hiBytes.castShape(floatSpecies, 0) as FloatVector +val inVecLo = FloatVector.fromArray(floatSpecies, input, inputStartLo + idx) +val inVecHi = FloatVector.fromArray(floatSpecies, input, inputStartHi + idx) +codeAccLo = inVecLo.fma(codeVecLo, codeAccLo) // for sub-block 2j +inputAccLo = inVecLo.add(inputAccLo) // dmin-correction sum +codeAccHi = inVecHi.fma(codeVecHi, codeAccHi) // for sub-block 2j+1 +inputAccHi = inVecHi.add(inputAccHi) +---- + +This halves the number of byte loads vs the prior helper that ran +once per nibble pass. Lives at +`skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel.kt`, +exposed via the `Q4KMatmulKernel` SPI sibling (kernel SPI: +`KernelProvider.matmulQ4K(): Q4KMatmulKernel?`). + +The MemSeg variant (`matmulF32Q4_KMemSeg`) is the same algorithm with +`ByteVector.fromMemorySegment` — no SPI surface yet, just an inline +replacement. + +==== Numbers + +`QuantizedMatmulBench`, JDK 21.0.10, Apple Silicon: + +[cols="1,1,1",options="header"] +|=== +| Shape (inputDim × outputDim) | Time | Throughput +| 1024 × 1024 | 0.07 ms | ~30 GFLOPS +| 4096 × 1024 | 0.15 ms | ~55 GFLOPS +| 4096 × 4096 | 0.46 ms | ~73 GFLOPS +|=== + +Same throughput regime as the FP32 SIMD kernel — the fused dequant +adds essentially zero cost on top of the FMA. + +=== Q6_K — 256 elements / 210 bytes / 16 sub-blocks + +Block layout: + +* bytes `[0, 128)`: `ql` — low 4 bits of each 6-bit code (half-interleaved) +* bytes `[128, 192)`: `qh` — high 2 bits of each 6-bit code (4 codes per byte) +* bytes `[192, 208)`: 16 signed `int8` sub-block scales +* bytes `[208, 210)`: FP16 `d` (super-block scale) + +Per element: 6-bit code = `(ql_nibble) | ((qh_2bits) << 4) − 32`, +dequant = `d · sc[sub_block] · code`. + +Q6_K's `qh` is the wrinkle: each `qh` byte carries the high 2 bits of +*four* codes, packed at bit positions 0–1, 2–3, 4–5, 6–7. The SIMD +recipe in `dequantQ6_KBlock` (file +`JvmQuantizedVectorKernels.kt`): + +[source,kotlin] +---- +val ql0Vec = ByteVector.fromArray(byteSpeciesForFloat, weight, qlBase + l) +val ql32Vec = ByteVector.fromArray(byteSpeciesForFloat, weight, qlBase + l + 32) +val qhVec = ByteVector.fromArray(byteSpeciesForFloat, weight, qhBase + l) + +val q1Bytes = ql0Vec.and(0x0F.toByte()) + .or(qhVec.and(0x03.toByte()).lanewise(LSHL, 4.toByte())) +val q2Bytes = ql32Vec.and(0x0F.toByte()) + .or(qhVec.lanewise(LSHR, 2.toByte()).and(0x03.toByte()).lanewise(LSHL, 4.toByte())) +val q3Bytes = ql0Vec.lanewise(LSHR, 4.toByte()) + .or(qhVec.lanewise(LSHR, 4.toByte()).and(0x03.toByte()).lanewise(LSHL, 4.toByte())) +val q4Bytes = ql32Vec.lanewise(LSHR, 4.toByte()) + .or(qhVec.lanewise(LSHR, 6.toByte()).lanewise(LSHL, 4.toByte())) +---- + +Then each `q*Bytes` is cast to `FloatVector`, biased by −32, scaled +by `d · sc[sub_block]`, and stored to one of four 32-element regions +of the per-block scratch. The per-cell SIMD dot product (`FloatVector.fma`) +already existed; this change only replaced the scalar dequant loop. + +=== Q4_0 — 32 elements / 18 bytes + +The simplest layout (single FP16 scale + 16 packed nibble bytes), but +also the *least* SIMD-friendly: adjacent elements share a byte +(`code[2k]` lo, `code[2k+1]` hi), so getting codes in natural element +order from a `ByteVector` would need a lane-interleave shuffle or a +strided gather. NEON has no native gather instruction, so on Apple +Silicon a gather-based pipeline would fall back to scalar loads. + +The shipped Q4_0 implementation uses the *partial-vec* pattern Q4_K +used before its fully-fused rewrite: + +[source,kotlin] +---- +// Stage 1: scalar byte-pair unpack into a 32-element scratch FloatArray. +for (k in 0 until 16) { + val b = weightSeg.get(JAVA_BYTE_LE, codesOffset + k.toLong()).toInt() and 0xFF + codeBuf[2 * k] = (b and 0x0F).toFloat() - 8f + codeBuf[2 * k + 1] = (b ushr 4).toFloat() - 8f +} +// Stage 2: SIMD FMA dot product. +var accVec = FloatVector.zero(floatSpecies) +while (idx < loopBound) { + val iv = FloatVector.fromArray(floatSpecies, input, inputOffset + idx) + val cv = FloatVector.fromArray(floatSpecies, codeBuf, idx) + accVec = iv.fma(cv, accVec) + idx += step +} +return (accVec.reduceLanes(ADD) + scalarTail) * scale +---- + +If Q4_0 ever becomes a hot path (it's rarely seen in modern weights — +Q4_K_M / Q4_K_S dominate Gemma 4, Llama, Qwen), the upgrade to a +fully-fused `ByteVector` pipeline is a reasonable follow-up — same +shape as the Q4_K rewrite, with the lane-interleave done via +`VectorShuffle`. + +== Per-format coverage matrix + +[cols="1,1,1,2",options="header"] +|=== +| Format | SPI sibling? | MemSeg variant SIMD? | Inner loop strategy +| Q8_0 | no | yes | Fully fused (`ByteVector.castShape` + scaled FMA) +| Q4_K | yes (`Q4KMatmulKernel`) | yes (inline, same algorithm) | Fully fused (single byte load → lo+hi nibble accumulators, lazy `dmin`) +| Q6_K | no | n/a | SIMD dequant into scratch + SIMD dot (two-stage) +| Q4_0 | no | yes | Scalar unpack into scratch + SIMD dot (two-stage) +|=== + +== Where to look in the code + +[cols="1,2",options="header"] +|=== +| File | What it covers +| `skainet-backends/skainet-backend-api/.../kernel/Q4KMatmulKernel.kt` | The Q4_K kernel SPI (commonMain). +| `skainet-backends/skainet-backend-cpu/src/jvmMain/.../kernel/PanamaVectorQ4KMatmulKernel.kt` | Fused-pipeline Q4_K implementation. +| `skainet-backends/skainet-backend-cpu/src/jvmMain/.../tensor/ops/JvmQuantizedVectorKernels.kt` | All the per-format kernels — Q4_K MemSeg, Q6_K dequant, Q4_0, Q8_0 MemSeg variants. +| `skainet-backends/skainet-backend-cpu/src/jvmMain/.../tensor/ops/DefaultCpuOpsJvm.kt` | `chooseQuantizedMatmul` — production dispatch by tensor data type. +| `skainet-backends/benchmarks/jvm-cpu-jmh/src/jmh/kotlin/sk/ainet/bench/QuantizedMatmulBench.kt` | JMH harness for Q4_K Panama. +| `skainet-backends/skainet-backend-cpu/src/jvmTest/.../kernel/PanamaVectorQ4KMatmulKernelTest.kt` | Parity tests vs the reference partial-vec kernel. +|=== + +For the kernel SPI itself, see +xref:explanation/perf/simd-kernels.adoc[]. For the planned native FFM +provider that would replace the Vector path on supported hosts, see +xref:explanation/perf/native-ffm-plan.adoc[]. diff --git a/docs/modules/ROOT/pages/explanation/perf/simd-kernels.adoc b/docs/modules/ROOT/pages/explanation/perf/simd-kernels.adoc new file mode 100644 index 000000000..b3d2a8e42 --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/perf/simd-kernels.adoc @@ -0,0 +1,244 @@ += How SIMD Kernels Are Built +:description: How SKaiNET wires JDK Vector API SIMD kernels behind a pluggable provider SPI for FP32 matmul. + +This page explains *how* SKaiNET's CPU backend reaches near-peak SIMD +throughput on the JVM, not *what* the public ops do. If you just want +to call `matmul` and have the fastest available kernel run, you don't +need to read this — `DirectCpuExecutionContext` does the right thing. +This page is for the engineer who wants to understand or extend the +kernel layer. + +== The kernel SPI in two paragraphs + +SKaiNET separates *what to compute* (`TensorOps.matmul`) from *which +SIMD recipe to compute it with* (`Fp32MatmulKernel.matmul`). The +separation lives in module `skainet-backend-api`: a tiny common-code +SPI of three types — `KernelProvider`, `Fp32MatmulKernel`, +`KernelRegistry` — that any backend can implement without depending on +the rest of SKaiNET. + +A *provider* is a bundle of kernels for one execution recipe (scalar, +Panama Vector, future native FFM, future GPU). Each provider declares +a `priority` and an `isAvailable()` check. At runtime, +`KernelRegistry.bestAvailable()` picks the highest-priority provider +that reports itself available, and the production op set +(`DefaultCpuOpsJvm.matmul`) pulls the kernel for the dtype it needs +(`matmulFp32()` returns `Fp32MatmulKernel?`, with `null` meaning "I +don't carry that"). Three providers ship with the CPU backend today: + +[cols="1,1,1,3",options="header"] +|=== +| Provider | Priority | When available | Notes +| `ScalarKernelProvider` | 0 | always | Three-loop reference; the parity baseline. +| `PanamaVectorKernelProvider` | 50 | JDK 21+ with `--add-modules jdk.incubator.vector` and `skainet.cpu.vector.enabled != false` | Tile-blocked FMA; the production winner on every supported JVM. +| (future) `NativeKernelProvider` | 100 | JDK 22+ with the native lib loaded | Captured as a plan in xref:explanation/perf/native-ffm-plan.adoc[]; not yet shipped. +|=== + +== Why the SPI exists + +Three reasons, in order of importance: + +. *Kernel-level benchmarking.* `KernelMatmulBench` in +`:skainet-backends:benchmarks:jvm-cpu-jmh` measures +`Fp32MatmulKernel.matmul` directly with no `TensorOps` wrapping — +the timed region is just the SIMD loop. Without an SPI we'd have to +either expose internals (leak production code) or duplicate the +algorithm in the bench (drift). +. *Drop-in replacement of the Vector path with native code* once a +hand-tuned NEON / AVX2 routine ever ships. The op layer doesn't +need to know. +. *Test isolation.* A scalar reference is always `null`-fallback-safe; +parity tests can pin it explicitly via +`ScalarKernelProvider.matmulFp32()` without going through the +registry. + +== JDK Vector API patterns used by `PanamaVectorMatmulKernel` + +The kernel lives at +`skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt` +and computes `C(m, n) = A(m, k) · B(k, n)` for row-major dense +tensors. Four ideas drive the implementation: + +=== 1. `FloatVector.SPECIES_PREFERRED` + +`SPECIES_PREFERRED` returns the widest float vector species the host +CPU can execute natively — 4 lanes on Apple Silicon NEON (128-bit), 8 +lanes on AVX2 (256-bit), 16 lanes on AVX-512 (512-bit). Code that +loops on `species.length()` therefore unrolls once per platform +without a recompile. + +[source,kotlin] +---- +private val species: VectorSpecies = FloatVector.SPECIES_PREFERRED +val step = species.length() +val loopBound = species.loopBound(k) // largest multiple of `step` ≤ k +---- + +`loopBound` is the safe vectorized loop limit; the *tail* (elements +that don't fit a full vector) is handled scalar after the SIMD loop. + +=== 2. Pack `B^T` once per call, stream contiguously + +The naive matmul has a strided access pattern over `B`: for each +output cell `C[i, j]` the inner reduction walks column `j` of `B`, +which jumps `n` floats per step. That's a cache-line miss per element. + +`PanamaVectorMatmulKernel` packs `B` into a transposed buffer `bt` of +shape `(n, k)` once per matmul: + +[source,kotlin] +---- +val bt = FloatArray(n * k) +for (kk in 0 until k) { + val src = bOffset + kk * bStride + for (j in 0 until n) bt[j * k + kk] = b[src + j] +} +---- + +Now `bt[j, *]` is contiguous over `k` and the inner reduction streams +`a[i, kk] * bt[j, kk]` for both operands sequentially. + +=== 3. FMA accumulator + lane-reduce + +The inner loop is a vector-width fused multiply-add into a vector +accumulator, reduced once per output cell: + +[source,kotlin] +---- +var acc = FloatVector.zero(species) +var idx = 0 +while (idx < bound) { + val va = FloatVector.fromArray(species, a, aRowBase + idx) + val vb = FloatVector.fromArray(species, bt, btRowBase + idx) + acc = va.fma(vb, acc) + idx += step +} +var sum = acc.reduceLanes(VectorOperators.ADD) +// scalar tail +while (idx < kLen) { sum += a[aRowBase + idx] * bt[btRowBase + idx]; idx++ } +out[outRow + j] = sum +---- + +`fma(vb, acc) = this · vb + acc` in a single hardware instruction, ~2 +FLOPs per cycle per lane on hardware that supports FMA (every modern +x86_64 and ARMv8). `reduceLanes(ADD)` collapses the 4/8/16-lane +accumulator to one float — typically a tree of pairwise adds inside +the JIT, no scalar fallback. + +=== 4. Cache blocking with 8×8×128 tiles + +The ideas above are SIMD over `k` (the contraction axis). The next +bottleneck is L1 reuse over `m × n`: at large sizes, each row of `bt` +that we stream through the inner loop falls out of L1 before the next +output row gets to reuse it. + +The fix is *tile blocking* — three nested outer loops over `(m, n, +k)`-tiles, with the SIMD loop inside the innermost tile: + +[source,kotlin] +---- +private const val TILE_M = 8 +private const val TILE_N = 8 +private const val TILE_K = 128 + +while (mTile < m) { + val mEnd = minOf(mTile + TILE_M, m) + while (nTile < n) { + val nEnd = minOf(nTile + TILE_N, n) + while (kTile < k) { + val kEnd = minOf(kTile + TILE_K, k) + // SIMD inner loop runs on the (8 × 8 × 128) sub-block. + // Output is zeroed once before all tiles run; the K-tile + // loop accumulates via `+=` so partial K-tiles add cleanly. + } + } +} +---- + +8×8×128 floats ≈ 8 KB working set per tile, comfortably inside any +modern L1. On Apple Silicon NEON this brings the FP32 SIMD kernel from +~118 ms (untiled, B^T pack only) to ~80 ms at 1024² — ~10× over +scalar, on par with the prior production blocked path. + +== Auto-discovery: ServiceLoader + factory wrappers + +The kernel registry doesn't auto-load anything by default — that lets +tests instantiate it empty and pin specific providers. For real +applications, `KernelServiceLoader.installAll()` (in module +`skainet-backend-api`, JVM source set only) does: + +[source,kotlin] +---- +fun installAll(): List { + val providers = ServiceLoader.load(KernelProvider::class.java).toList() + for (p in providers) KernelRegistry.register(p) + return providers.map { it.name } +} +---- + +Because `ServiceLoader` requires a public no-arg constructor and +Kotlin `object` declarations don't expose one, the CPU backend ships +*factory wrappers* whose only job is to delegate to the singletons: + +[source,kotlin] +---- +public class ScalarKernelProviderFactory : KernelProvider by ScalarKernelProvider +public class PanamaVectorKernelProviderFactory : KernelProvider by PanamaVectorKernelProvider +---- + +Both classes are listed in +`skainet-backends/skainet-backend-cpu/src/jvmMain/resources/META-INF/services/sk.ainet.backend.api.kernel.KernelProvider`, +which is what `ServiceLoader.load` scans. + +`DefaultCpuOpsJvm.matmul` triggers `installAll()` lazily on first use +when the registry is empty, so callers don't need to wire it +themselves: + +[source,kotlin] +---- +private val fp32MatmulKernel: Fp32MatmulKernel by lazy { + if (KernelRegistry.providers().isEmpty()) { + KernelServiceLoader.installAll() + } + KernelRegistry.bestAvailable()?.matmulFp32() ?: ScalarMatmulKernel +} +---- + +== Numbers + +`KernelMatmulBench` in `:skainet-backends:benchmarks:jvm-cpu-jmh`, +JDK 21.0.10 on Apple Silicon, 3 warmup × 5 measurement iterations: + +[cols="1,1,1,1",options="header"] +|=== +| Size | Scalar | Panama (tiled) | Speedup +| 256² | 9.77 ms | 1.13 ms | *8.61×* +| 512² | 81.55 ms | 9.47 ms | *8.62×* +| 1024² | 865.54 ms | 79.88 ms | *10.83×* +|=== + +Compared to the prior production path (`JvmVectorKernels.matmulFloatBlocked`, +which was *not* behind the SPI), the SPI tiled kernel is 8.5% +faster at 256², 8.8% faster at 512², and within JMH noise at 1024² — +so wiring `DefaultCpuOpsJvm.matmul` to go through the registry was a +no-regression change. + +== Where to look in the code + +[cols="1,2",options="header"] +|=== +| File | What it does +| `skainet-backends/skainet-backend-api/.../kernel/KernelProvider.kt` | The SPI interface (per-kernel accessors, `priority`, `isAvailable()`). +| `skainet-backends/skainet-backend-api/.../kernel/KernelRegistry.kt` | Process-wide registry; `bestAvailable()` lookup. +| `skainet-backends/skainet-backend-api/src/jvmMain/.../KernelServiceLoader.kt` | JVM-only auto-discovery via `ServiceLoader`. +| `skainet-backends/skainet-backend-cpu/.../kernel/ScalarMatmulKernel.kt` | Three-loop reference (commonMain, all targets). +| `skainet-backends/skainet-backend-cpu/src/jvmMain/.../kernel/PanamaVectorMatmulKernel.kt` | The tile-blocked Vector-API implementation. +| `skainet-backends/skainet-backend-cpu/src/jvmMain/.../tensor/ops/DefaultCpuOpsJvm.kt` | Production routing — `matmul` resolves the SPI kernel lazily. +| `skainet-backends/benchmarks/jvm-cpu-jmh/src/jmh/kotlin/sk/ainet/bench/KernelMatmulBench.kt` | Direct kernel-level JMH harness used to capture the numbers above. +|=== + +For quantized matmul (Q4_K, Q6_K, Q8_0, Q4_0) — same story, different +inner loop — see xref:explanation/perf/quantized-simd-kernels.adoc[]. + +For the still-unbuilt native FFM provider, see +xref:explanation/perf/native-ffm-plan.adoc[]. diff --git a/docs/modules/ROOT/pages/reference/architecture.adoc b/docs/modules/ROOT/pages/reference/architecture.adoc index d350b26fc..21392ebce 100644 --- a/docs/modules/ROOT/pages/reference/architecture.adoc +++ b/docs/modules/ROOT/pages/reference/architecture.adoc @@ -1,11 +1,288 @@ = Architecture -:description: How SKaiNET's compile and execution layers are organized. +:description: Arc42-style architecture reference for SKaiNET — solution strategy, building blocks, runtime view, with focus on the kernel SPI and eager-execution pipeline. -SKaiNET uses a hybrid backend strategy that separates development -iteration from production deployment. +This page follows arc42's chapter ordering at a coarse grain. Sections +that are still single-paragraph stubs ask for contributions; sections +about the kernel SPI and eager execution pipeline are the deepest +because that's where the 0.21.0 work landed. + +== 1. Introduction and goals + +SKaiNET is a Kotlin Multiplatform ML framework whose primary target is +*on-device / edge inference and training* in environments that already +have JVM tooling: Android, JVM server, and Kotlin/Native iOS and +linuxX64. The framework separates *model authoring* (a typed Kotlin +DSL with compile-time tensor shape checks where possible) from +*execution* (pluggable backends), so the same model can run eagerly +during development and be lowered to MLIR StableHLO → IREE for +deployment. + +Hard non-goals: + +* *Become a full numerics library.* SKaiNET targets the operators +real models use, not the long tail in PyTorch / NumPy. +* *Run untrusted user code.* Kernels are trusted code; security is +about not corrupting memory, not about sandboxing. image::SKaiNET-compiler.svg[Architecture diagram of the SKaiNET compiler pipeline] -// The original ARCHITECTURE.md at the repo root was a 4-line stub -// pointing at the compiler diagram. If you are looking for a -// deeper architecture write-up, contribute it as a PR to this page. +== 2. Constraints + +[cols="1,3",options="header"] +|=== +| Constraint | Why +| Kotlin Multiplatform with `commonMain` / per-target source sets | Same DSL must run on JVM, Android, iOS, macOS, linuxX64, JS, Wasm. +| `--enable-preview --add-modules jdk.incubator.vector` on JVM 21+ | FloatVector / ByteVector are still incubator on JDK 25 (JEP 508). +| Maven Central publication via `vanniktech.mavenPublish` | All modules signed; coordinates `sk.ainet.core:*`. +| Antora-based docs site under `docs/modules/ROOT/` | Source-controlled, follows Diátaxis quadrants for user-facing pages. +| arc42 ordering for *this* page | Architectural reference, not a tutorial. +|=== + +== 3. Context (system boundaries) + +The framework's outer surface: + +* *DSL layer* — Kotlin model DSL (`nn { ... }`, `tensor { ... }`) and +imperative `TensorOps` / `ExecutionContext` API. +* *I/O layer* — model loaders for GGUF, SafeTensors, ONNX (read-only) +in `skainet-io-*` modules. +* *Compile layer* — `RecordingExecution` records ops to a tape, then +lowers to StableHLO / IREE bytecode in `skainet-compile-*` modules. +* *Backend layer* — `BackendProvider` dispatches `TensorOps` calls +to a concrete implementation (CPU, XNNPACK, future GPU). Inside a +backend, the *kernel SPI* picks the SIMD recipe for the host +hardware. + +== 4. Solution strategy + +SKaiNET runs the *same model graph* through one of two execution +strategies: + +* *Eager execution* — `DirectCpuExecutionContext` calls op +implementations as the user invokes them. Used during development, +testing, and on-device inference paths where AOT compilation is +impractical (debug builds, dynamic graphs). This is the path the +0.21.0 SIMD work targets. +* *Recorded execution* — `RecordingExecution` builds an op tape, +which `HloGenerator` lowers to StableHLO MLIR. IREE compiles the +MLIR to a portable bytecode for production deployment. + +Both strategies share the same `TensorOps` surface, so a model +written once runs in either mode without changes. Numerical parity +between modes is part of the test contract. + +== 5. Building block view (static structure) + +=== 5.1 Module layout + +[cols="2,3",options="header"] +|=== +| Module path | Role +| `skainet-lang/skainet-lang-core` | DSL types, tensor abstractions, common ops, `TensorOps` / `ExecutionContext` interfaces. KMP, all targets. +| `skainet-lang/skainet-lang-models` | Reference reusable models (Llama, Gemma, Qwen, Whisper) built on the DSL. +| `skainet-backends/skainet-backend-api` | Neutral backend SPI — `TensorOps`, `TensorDataFactory`, **kernel SPI** (`KernelProvider`, `Fp32MatmulKernel`, `Q4KMatmulKernel`, `KernelRegistry`). +| `skainet-backends/skainet-backend-cpu` | CPU implementation. Eager-execution `DefaultCpuOpsBase` (commonMain) + `DefaultCpuOpsJvm` (jvmMain) with SIMD kernels. +| `skainet-backends/skainet-backend-xnnpack` | Optional XNNPACK CPU backend (FP32 matmul / conv2d / pooling) on linuxX64 / linuxArm64 / Android. +| `skainet-backends/benchmarks/jvm-cpu-jmh` | JMH harness — `MatmulBench`, `KernelMatmulBench`, `QuantizedMatmulBench`, `ElementwiseAdd1MBench`, `Reductions1MBench`. +| `skainet-compile/*` | Tape recording, StableHLO emission, IREE export. +| `skainet-io/*` | Model loaders (GGUF, SafeTensors, ONNX), tokenizers, IRPA writer. +|=== + +=== 5.2 Kernel SPI + +Introduced in 0.21.0 (PRs #554, #559, #562). The static structure: + +[source] +---- + commonMain (skainet-backend-api) + ┌──────────────────────────────────────┐ + │ KernelProvider { │ + │ name: String │ + │ priority: Int │ + │ isAvailable(): Boolean │ + │ matmulFp32(): Fp32MatmulKernel? │ + │ matmulQ4K(): Q4KMatmulKernel? │ + │ } │ + │ │ + │ KernelRegistry { │ + │ register(KernelProvider) │ + │ bestAvailable(): KernelProvider? │ + │ find(name): KernelProvider? │ + │ } │ + │ │ + │ Fp32MatmulKernel.matmul(...) │ + │ Q4KMatmulKernel.matmul(...) │ + └──────────────┬───────────────────────┘ + │ implements / extends + ┌────────────────────────────┼────────────────────────────────┐ + │ jvmMain (api) │ commonMain (cpu) │ + │ KernelServiceLoader │ ScalarMatmulKernel (priority 0) │ + │ installAll() │ ScalarKernelProvider │ + └────────────────────────────┴────────────────────────────────┘ + │ + ┌────────────────────────────┼────────────────────────────────┐ + │ jvmMain (cpu) │ + │ PanamaVectorMatmulKernel (priority 50, FP32) │ + │ PanamaVectorQ4KMatmulKernel (priority 50, Q4_K) │ + │ PanamaVectorKernelProvider │ + │ Scalar/PanamaVectorKernelProviderFactory (no-arg wrappers) │ + │ META-INF/services/...KernelProvider │ + └─────────────────────────────────────────────────────────────┘ +---- + +Three live providers ship; a fourth (priority 100, native FFM) is +captured as a plan in xref:explanation/perf/native-ffm-plan.adoc[] +and not yet built. For *how* the kernels are implemented, see +xref:explanation/perf/simd-kernels.adoc[] (FP32) and +xref:explanation/perf/quantized-simd-kernels.adoc[] (quantized). + +== 6. Runtime view — eager execution + +The eager pipeline for a single op: + +[source] +---- +User code + │ + │ ctx.ops.matmul(a, b) + ▼ +DirectCpuExecutionContext + │ + │ delegates to TensorOps implementation + ▼ +DefaultCpuOpsJvm.matmul (jvmMain) + │ + ├─→ chooseQuantizedMatmul(a, b)? + │ │ + │ │ matches Q4_K / Q6_K / Q8_0 / Q4_0 weights + │ ▼ + │ ┌─ Q4_K branch: q4kMatmulKernel?.matmul(...) via SPI + │ │ fallback to JvmQuantizedVectorKernels.matmulQ4_KVec + │ ├─ Q6_K branch: JvmQuantizedVectorKernels.matmulQ6_KVec + │ ├─ Q8_0 / Q4_0 branches: per-format SIMD inner loops + │ └─ MemSeg variants: same kernels, ByteVector.fromMemorySegment + │ + ├─→ chooseMatmul(a, b)? (FP32 path) + │ │ + │ │ fp32MatmulKernel.matmul(...) ← SPI dispatch + │ ▼ + │ PanamaVectorMatmulKernel.matmul(...) + │ │ (or ScalarMatmulKernel when Panama unavailable) + │ │ + │ │ tile-blocked FMA inner loop: + │ │ for each (m, n, k)-tile: + │ │ load FloatVector slices of A and B^T + │ │ acc = va.fma(vb, acc) + │ │ acc.reduceLanes(ADD) per output cell + │ ▼ + │ FloatArray output ← back up the stack + │ + └─→ super.matmul(a, b) (DefaultCpuOpsBase fallback) +---- + +Two specifics worth calling out: + +* *Lazy provider resolution.* `DefaultCpuOpsJvm.fp32MatmulKernel` and +`q4kMatmulKernel` are `by lazy` properties. First access triggers +`KernelServiceLoader.installAll()` if the registry is empty, then +caches the resolved kernel for the lifetime of the op set. Apps +that pre-register custom providers via +`KernelRegistry.register(...)` before constructing the op set bypass +the auto-discovery path. +* *Fall-through everywhere.* Each routing decision (`chooseQuantizedMatmul` +→ `chooseMatmul` → `super.matmul`) returns `null` on a miss, never +throws, so adding a new tensor type or a new SPI accessor is purely +additive. The MemSeg arena leak fix in PR #556 made every per-op +output use `Arena.ofAuto()`, so even the fast-path branches don't +need explicit lifetime management. + +For shape inference, broadcasting, and the lazy shape-swap transpose +specializations on `Q4_KTensorData` / `Q6_KTensorData` / `MemorySegmentBackedData`, +see the same source file (`DefaultCpuOpsJvm.transpose`). + +== 7. Deployment view + +* *Maven Central* — every module published as +`sk.ainet.core:-:`. KMP variants land per +target (`*-jvm`, `*-android`, `*-iosarm64`, `*-macosarm64`, +`*-linuxx64`, `*-linuxarm64`, `*-js`, `*-wasm-js`, `*-wasm-wasi`). +* *Single BOM* — `sk.ainet.core:skainet-bom` provides a +`platform()` import for downstream Gradle. +* *Releases* — tags `0.X.Y` on the release branch trigger +`.github/workflows/publish.yml` → `./gradlew publish` on macOS-latest +with JDK 25. + +== 8. Cross-cutting concepts + +* *Numerical parity testing.* Every accelerated kernel has a parity +test against a scalar reference within a documented tolerance +(typically `1e-5 * k` or `1e-4` relative). Examples: +`PanamaVectorMatmulKernelTest`, `PanamaVectorQ4KMatmulKernelTest`, +`Q6KMatmulTest`. The scalar reference is the contract; SIMD speed +is a non-functional improvement that must not break the contract. +* *Lazy resource lifetimes.* `MemorySegmentTensorDataFactory` uses +`Arena.ofAuto()` for per-op outputs so output segments are +GC-reclaimable. The earlier `Arena.ofConfined()` builds leaked +~tens of MB per matmul, blowing 32+ GiB of direct memory in +inference loops over a 35-layer Gemma 4 forward pass. Fixed in +PR #556. +* *Kill switches via system properties.* The Vector API code path +respects `-Dskainet.cpu.vector.enabled=false` so a deployment can +opt out of incubator code without a recompile. Same pattern for +BLAS (`-Dskainet.cpu.blas.enabled=true`). + +== 9. Architecture decisions + +[cols="1,1,3",options="header"] +|=== +| Decision | Date | Rationale +| Kernel SPI parallel to BackendProvider | 2026-04 (PR #554) | Matmul / SDPA are model-agnostic; isolating them lets bench harnesses time the SIMD loop directly and lets a future native provider register without touching the op layer. +| `KernelProvider.matmulQ4K()` accessor with `default null` | 2026-04 (PR #562) | Backwards compat for existing providers (Scalar) without forcing every implementation to override. Same pattern will be used for Q6KMatmulKernel / Q4KMemSegMatmulKernel sibling SPIs. +| ServiceLoader auto-discovery deferred until 2 providers exist | 2026-04 (PR #559) | Single-provider auto-discovery would have been ceremony for nothing; once Panama landed alongside Scalar, the trigger condition was met. +| FFM (not JNI) for any future native code | roadmap M5 | JNI's per-call overhead and global lock are wrong for hot per-token kernels. +| Antora docs (Diátaxis), not GitHub Wiki | 2025 | Source-controlled, branchable, ranked higher in search than wikis, ships with the repo. +|=== + +== 10. Quality requirements + +* *Performance.* Panama FP32 matmul ≥1.5× scalar (M5 metric — *met*, +~10× at 1024² on Apple Silicon NEON). Native Q4_K matmul ≥2.5× +scalar dequant baseline (M5 metric — *deferred to native FFM PR*). +* *Numerical equivalence.* Every SIMD kernel matches its scalar +reference within FP-rounding tolerance (`1e-5 * k` for matmul, +`1e-4` relative for quantized). Pinned by parity tests. +* *Multi-target buildability.* `./gradlew allTests` (the release +gate) must pass on every KMP target — JVM, JS, Wasm, macosArm64, +iosSimulatorArm64, linuxX64, linuxArm64, Android. + +== 11. Risks and technical debt + +* *Vector API still incubator on JDK 25.* JEP 508 keeps it that way +through 2026; we depend on it heavily. If it breaks API in a future +JDK, every `JvmVectorKernels` / `JvmQuantizedVectorKernels` file +needs adjustment. Mitigation: thin wrappers, parity tests are a +canary. +* *No native FFM provider yet.* The literal M5 milestone metric +(`≥2.5×` for Q4_K) is met by Panama in absolute terms but not in +the "native vs JVM" framing the metric originally specified. +Mitigation: `xref:explanation/perf/native-ffm-plan.adoc[]` +documents what shipping it would look like. +* *Two reverted optimizations on develop history.* MemSeg pool +(commit 8642b322) and intra-op matmul parallelism (commit +9ed633b6) were both tried and reverted. Re-attempts need a +different strategy than what was tried; the revert commits +explain why. + +== 12. Glossary (selected) + +[cols="1,3",options="header"] +|=== +| Term | Meaning +| FFM | Foreign Function & Memory API (JEP 442 et seq.). Java 22 stable, 21 preview. Replaces JNI for native interop. +| FMA | Fused multiply-add (`a · b + c` in one instruction). Supported by every modern x86_64 (FMA3) and ARMv8 CPU. +| ggml | The C library underpinning llama.cpp; defines the canonical Q4_K / Q6_K / Q8_0 block layouts SKaiNET uses for GGUF compatibility. +| Lazy transpose | Specialization in `DefaultCpuOpsJvm.transpose` that swaps a tensor's *shape* without reordering its packed bytes — works because the matmul kernel's byte-offset math is symmetric across the swap. +| MemSeg | Short for `java.lang.foreign.MemorySegment`. Off-heap memory abstraction used for mmap'd weight buffers. +| Panama | Codename for the JDK Vector API (`jdk.incubator.vector`) and FFM. Both originate from Project Panama. +| SPI | Service provider interface — a public interface with multiple registered implementations, looked up at runtime. SKaiNET uses it for backends and now for kernels. +|=== From 75aa4f6295d0bf22548c9777dfd8bdfa5c087eed Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 29 Apr 2026 07:48:09 +0200 Subject: [PATCH 2/2] docs: fix asciidoctor list-index warning in native-ffm-plan Local Antora build (the same docker pipeline GitHub Actions runs) emitted: warn: native-ffm-plan.adoc:237: list item index: expected 1, got 22 The line started with "22." after a wrap, which the asciidoctor parser interpreted as a sibling numbered list item with an out-of- sequence index. Re-wrap so "22." stays mid-sentence. Site rebuild now warning-free. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc b/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc index 3e25e01b8..9553f1efb 100644 --- a/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc +++ b/docs/modules/ROOT/pages/explanation/perf/native-ffm-plan.adoc @@ -233,8 +233,8 @@ the same numbers as today. == Risks & open questions -. *JDK 21 preview FFM vs JDK 22 stable.* FFM left preview in Java -22. The repo currently builds on JDK 21 with `--enable-preview +. *JDK 21 preview FFM vs JDK 22 stable.* FFM left preview in Java 22. +The repo currently builds on JDK 21 with `--enable-preview --add-modules jdk.incubator.vector`. Recommendation: stay on 21 preview; flip to 22 in a separate toolchain-bump PR. . *`MethodHandle` invocation overhead.* Even with FFM, each native