diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index 1723bb7c3..d0bcbd969 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -38,6 +38,7 @@ ** xref:explanation/virtual-tensors.adoc[Virtual tensors: one logical tensor, many physical forms] ** xref:explanation/packed-weight-layout.adoc[Packed weight layout] ** xref:explanation/eager-execution.adoc[Eager execution: backends and kernels] +** xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm] ** xref:explanation/quantization-process.adoc[The quantization process] ** xref:explanation/theory/index.adoc[Mathematical theory] *** xref:explanation/theory/matmul.adoc[Matrix multiplication] diff --git a/docs/modules/ROOT/pages/explanation/eager-execution.adoc b/docs/modules/ROOT/pages/explanation/eager-execution.adoc index 7cd46060a..a81b6a610 100644 --- a/docs/modules/ROOT/pages/explanation/eager-execution.adoc +++ b/docs/modules/ROOT/pages/explanation/eager-execution.adoc @@ -6,6 +6,11 @@ distinct from the StableHLO/IREE export path. This page is the hand-authored ove companion xref:reference/kernel-support-matrix.adoc[kernel × platform matrix] is generated from the registered providers and gated against drift. +This page covers the provider tier — who can compute what, on which platform. The generic path +selects differently, on a declared descriptor of the operands rather than on dtype and priority: +see xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm] for +`KernelDispatch`, the `ViewKernelPack` SPI, and how a call site is matched to a kernel. + Legend: ✅ available · ❌ missing. [mermaid] diff --git a/docs/modules/ROOT/pages/explanation/kernel-selection.adoc b/docs/modules/ROOT/pages/explanation/kernel-selection.adoc new file mode 100644 index 000000000..f4eae80cd --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/kernel-selection.adoc @@ -0,0 +1,214 @@ += Kernel SPI and the selection algorithm +:description: The two kernel registries, how a backend registers kernels, how a call site is matched to one, and why a miss is silent unless you look for it. + +A matmul in SKaiNET does not pick its implementation from an `is`-ladder over Kotlin classes. +It is *selected*, from a declared descriptor of the operands, out of a registry a backend +populated. This page explains the two registries, the SPI a backend implements, the exact +selection algorithm, and how to tell which kernel actually ran. + +Its companion pages: xref:explanation/eager-execution.adoc[Eager execution] maps the backends and +platform coverage; xref:explanation/packed-weight-layout.adoc[Packed weight layout] explains the +block orders this page matches on; xref:reference/kernel-support-matrix.adoc[the support matrix] is +generated from real registrations. + +== Two registries, two jobs + +There are two, and confusing them is the most common source of "why is this slow". + +[cols="1,3,3"] +|=== +| | `KernelRegistry` | `KernelDispatch` + +| Selects on +| dtype + provider priority +| a `KernelKey` describing *every* operand + +| SPI +| `KernelProvider` +| `ViewKernelPack` (installs `ViewKernel` s) + +| Answers +| "who provides the best FP32 GEMM here?" +| "which kernel takes *this* activation and *this* weight, in these layouts?" + +| Used by +| the legacy fast paths in `DefaultCpuOps*` (`chooseQuantizedMatmulHeap` and friends) +| `KernelDispatch.matmul`, the generic path (SKEEP-003 §5.1) + +| Priority model +| scalar 0 · Panama 50 · native FFM/JNI 100 — highest available wins +| exact key match; later registration wins for the same key +|=== + +`KernelRegistry` answers a question about *capability*; `KernelDispatch` answers a question about +*applicability*. A provider can be the best available and still be unable to take a particular +weight — off-heap, block-order-mismatched, strided — which is exactly what the key encodes. + +== The descriptor: `KernelKey` + +`KernelKey.matmul(a, b)` builds `("matmul", [OperandKey.of(a), OperandKey.of(b)], HOST, capabilities)`. +Each `OperandKey` carries the operand's `Format` (dtype + encoding, e.g. dense FP32, `Q4_K`) and its +`LayoutClass`, derived from the view's layout: + +`CONTIGUOUS`:: dense, unit-stride. +`STRIDED`:: dense, non-unit stride — what a transposed weight view looks like. +`BLOCKED_ROW_MAJOR`:: quantization blocks laid out along rows — the canonical GGUF order, as loaded. +`BLOCKED_INPUT_MAJOR`:: blocks laid out input-major — what the packed SIMD kernels read. + +[IMPORTANT] +==== +Key equality is **exact**. There is no subsumption and no fuzzy match: a kernel registered for +`BLOCKED_INPUT_MAJOR` is invisible to a lookup for `BLOCKED_ROW_MAJOR`, and because `capabilities` +is part of the data class, a kernel registered with capabilities that the lookup does not request +can never be found by it. This is deliberate — selection is meant to be a table lookup you can +reason about — but it means registering a kernel is not the same as it being reachable. +==== + +== The SPI a backend implements + +Two interfaces, in `skainet-backend-api`: + +`KernelProvider`:: a compute backend (scalar, Panama Vector, native FFM, JNI NEON, Accelerate). +Exposes `matmulFp32()`, the packed-quant entry points, `isAvailable()` and a priority. + +`ViewKernelPack`:: an installable set of `ViewKernel` s for `KernelDispatch`. One method, +`install()`, which must be idempotent and must register *nothing* rather than throw when its +platform support is absent (a missing native library, no vector unit). + +=== Registration and discovery + +[cols="1,2,2"] +|=== +| Platform | `KernelProvider` | `ViewKernelPack` + +| JVM +| `ServiceLoader` via `KernelServiceLoader.installAll()` +| `ServiceLoader`, discovered by `KernelDispatch.ensureInstalled()` + +| Android +| `ServiceLoader` (keep `META-INF/services` through packaging) +| `ServiceLoader`, same caveat + +| Kotlin/Native, wasm, JS +| manual — e.g. `installNativeKernels()` +| manual — call the pack's `install()` yourself +|=== + +A backend module declares its service the usual way, e.g. `skainet-backend-native-cpu` ships +`META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack` naming +`FfmRowMajorKernelPackFactory`. `ServiceLoader` needs a public no-arg constructor, which a Kotlin +`object` does not expose, so each pack ships a thin factory class delegating to the singleton — +the same shape `NativeKernelProviderFactory` has always used for providers. + +== Bootstrap: the dispatcher heals itself + +`KernelDispatch.matmul` calls `ensureInstalled()` first. When the table is empty and nothing has +been registered, it performs a one-time bootstrap: + +. discover `KernelProvider` s, if `KernelRegistry` is empty; +. `KernelPacks.install()`, which registers the reference kernel plus the best available provider's + dense-FP32 view kernels (contiguous *and* strided) and its `BLOCKED_INPUT_MAJOR` packed kernels; +. install every discovered `ViewKernelPack` — on JVM that is the FFM row-major pack, i.e. the + `BLOCKED_ROW_MAJOR` kernels that serve mapped GGUF weights zero-copy. + +[NOTE] +==== +Step 1 must come first. `KernelPacks.install()` defaults its provider to +`KernelRegistry.bestAvailable()`, which is `null` on an empty registry — bootstrapping in the wrong +order silently installs the reference kernel and nothing else. Measured on JVM: 8 of 17 kernels +land instead of all 17, and because the row-major pack installs unconditionally, a GGUF decode path +still looks fine while every dense-FP32 and input-block-major dispatch quietly runs on the +reference kernel. +==== + +Explicit registration still wins: a consumer that registers kernels before the first dispatch +suppresses auto-install entirely, so a curated set is never silently widened. `clearForTesting()` +re-arms the bootstrap. + +== The selection algorithm + +[mermaid] +---- +flowchart TD + A["matmul(a, b, out)"] --> B["ensureInstalled()
bootstrap if the table is empty"] + B --> C["normalizeActivation(a)
rank 1 → [1, k]; [b, s, k] → [b*s, k]"] + C --> D["key = KernelKey.matmul(a, b)"] + D --> E{"exact match?"} + E -->|yes| F["run it"] + E -->|no| G{"weight requests
another activation format?"} + G -->|"yes, e.g. ternary wants int8"| H["requantize into caller's Scope
emit AdapterInserted"] + H --> I{"kernel for the
requantized pair?"} + I -->|yes| F + G -->|no| J{"prepackWeights = true
and weight is ROW_MAJOR?"} + I -->|no| J + J -->|yes| K["prepack to INPUT_BLOCK_MAJOR
O(bytes), opt-in only"] + K --> L{"packed kernel now?"} + L -->|yes| F + J -->|no| M["adapt activation ('gather')"] + L -->|no| M + M --> N["ReferenceMatmulKernel
decodes any format, ~1000x slower"] +---- + +Two properties are worth stating plainly: + +* **Rank is normalised once, as views.** A rank-1 decode step never reaches a kernel written for + rank 2 — that class of `ClassCastException` disappears by construction. +* **Adapters are visible and caller-scoped.** When an operand must be converted, the allocation + happens in the caller's `Scope` (a `Forward` scope inside a generation loop) and is emitted as + `TraceEvent.AdapterInserted`, rather than hidden inside a kernel. + +`prepackWeights` is off by default on purpose: the relayout is O(bytes), so doing it inside a +decode step copies the whole weight *per token*. Prepack once at load instead +(`TensorView.prepack`), which then hits the exact key and copies nothing. + +== Why a kernel declines + +Reaching a kernel is not the same as it accepting the work. A kernel that cannot serve an operand +falls back and traces `reference-fallback from : `. The common reasons: + +* `FfmRowMajorMatmulKernel` — activation must be heap-backed `FloatArray` and contiguous; the + output must be a heap `FloatArray`; the weight must be buffer-backed or a heap `ByteArray`. +* `PackedViewMatmulKernel` — every operand must be `Storage.Heap`; off-heap and mapped storage are + not served by this tier. +* `Fp32ViewMatmulKernel` — the output must be a heap `FloatArray`, and the weight's row stride must + agree with its declared layout. +* Block alignment — a quantized tensor whose last dimension is not a multiple of the block size + (256 for K-quants, 32 for `Q4_0`/`Q8_0`) is rejected outright. + +Only `FfmRowMajorMatmulKernel` and `JniRowMajorMatmulKernel` implement `MappedCapableKernel`, i.e. +only they read a weight straight out of mapped or direct-buffer storage. That is why +`KernelDispatch.mappedServableEncodings()` is derived from live registrations rather than a +hand-kept list. + +== Diagnosing a selection + +The failure mode this design has to defend against is silence: the reference kernel is *correct* +for every format, so a miss produces right answers slowly rather than an error. + +`KernelDispatch.kernels()`:: what is actually registered, most recent first. On a healthy JVM +bootstrap this is 17 entries: 7 `ffm-rowmajor-*`, 7 `native-ffm-*` packed, 2 `native-ffm-fp32` +(contiguous and strided keys), and `reference`. +`KernelDispatch.mappedServableEncodings()`:: which encodings can be served zero-copy from a mapping +right now. +The one-time warning:: the first time the reference kernel serves a *blocked* weight, the +dispatcher prints what it could not match and how to install a pack. It fires once per process. +`KernelDispatch.defaultSink`:: set a real `TraceSink` to see every kernel run and adapter. +Production call sites (`DefaultCpuOps`) do not thread a sink through, so this global is how you +observe them. +`DispatchMode.useRegistry()`:: `-Dskainet.dispatch.registry=false` forces the legacy generic +fallback, which is useful for bisecting a suspected dispatch problem. + +== Adding a backend + +. Implement `KernelProvider`; add a no-arg factory class; list it in + `META-INF/services/sk.ainet.backend.api.kernel.KernelProvider`. +. If the backend has kernels that read a specific *layout* (packed, mapped, prepacked), implement + `ViewKernelPack`, add its factory to + `META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack`, and register one `ViewKernel` per + `(encoding, layout)` you actually serve. Register nothing for the rest — the reference kernel + keeps those correct. +. Make `install()` a no-op when the platform cannot support it, so discovery on a machine without + your native library costs a lookup and changes nothing. +. On Kotlin/Native, wasm and JS, document the manual install call — there is no discovery there. +. Regenerate the xref:reference/kernel-support-matrix.adoc[support matrix]; it is gated against + drift, so a new tier that forgets this fails the build. diff --git a/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.android.kt b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.android.kt new file mode 100644 index 000000000..df7edc218 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.android.kt @@ -0,0 +1,30 @@ +package sk.ainet.backend.api.kernel + +import java.util.ServiceLoader +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * Android discovery for [ViewKernelPack]. `ServiceLoader` exists on Android, so the JNI packs a + * consumer ships (e.g. the NEON row-major pack in `skainet-backend-jni-cpu`) are discovered the + * same way as on the JVM, provided the packaging step keeps `META-INF/services` entries. + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelPacks(): List = + runCatching { + ServiceLoader.load(ViewKernelPack::class.java) + .mapNotNull { pack -> runCatching { pack.install(); pack.name }.getOrNull() } + .toList() + }.getOrElse { emptyList() } + +/** + * Provider discovery, inlined rather than delegated to `KernelServiceLoader`: that object lives in + * `jvmMain`, which the Android source set does not see. Same two steps it performs — discover, then + * register, letting [KernelRegistry] sort by priority on insertion. + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = + runCatching { + ServiceLoader.load(KernelProvider::class.java) + .mapNotNull { provider -> runCatching { KernelRegistry.register(provider); provider.name }.getOrNull() } + .toList() + }.getOrElse { emptyList() } diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt index 548f8f115..b7c4694ac 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt @@ -28,6 +28,35 @@ public object KernelDispatch { private val kernels: MutableList = mutableListOf() + private var autoInstallAttempted: Boolean = false + + /** + * Populate the table from platform-discovered providers and [ViewKernelPack]s, once per + * process, when nothing has been registered yet. + * + * [KernelRegistry] has always self-healed this way (`DefaultCpuOpsJvm.ensureKernelProviders` + * installs providers on first use); this dispatcher did not, so every consumer had to remember + * an explicit bootstrap before its first forward pass. Forgetting it is silent — dispatch + * simply falls to the decoding reference kernel, which is correct and about a thousand times + * slower — and it was forgotten repeatedly in practice, by application entry points and + * diagnostic harnesses alike. + * + * Order matters: providers first, because [KernelPacks.install] derives its kernels from + * `KernelRegistry.bestAvailable()` and would otherwise contribute nothing but the reference + * kernel. + * + * Explicit installation still works and still wins — a consumer that registers its own kernels + * before the first dispatch suppresses auto-install entirely, and later registrations override + * earlier ones for the same key. Call [clearForTesting] to re-arm. + */ + public fun ensureInstalled() { + if (autoInstallAttempted || kernels.isNotEmpty()) return + autoInstallAttempted = true + if (KernelRegistry.providers().isEmpty()) installPlatformKernelProviders() + KernelPacks.install() + installPlatformKernelPacks() + } + /** Register [kernel]; later registrations win for the same key (a pack can override the reference). */ public fun register(kernel: ViewKernel) { kernels.removeAll { it.key == kernel.key && it.name == kernel.name } @@ -40,7 +69,10 @@ public object KernelDispatch { /** The kernel registered for [key], or `null`. */ public fun find(key: KernelKey): ViewKernel? = kernels.firstOrNull { it.key == key } - public fun clearForTesting() { kernels.clear() } + public fun clearForTesting() { + kernels.clear() + autoInstallAttempted = false + } /** * Encodings a [MappedCapableKernel] registered right now serves as a `BLOCKED_ROW_MAJOR` @@ -76,6 +108,16 @@ public object KernelDispatch { } } + /** + * Process-global default [TraceSink] used when a call site does not pass one. Production + * call sites (e.g. `DefaultCpuOps`) rely on the parameter default, which made every + * reference-kernel fallback invisible — set this (e.g. from a diagnostic harness) to + * observe dispatch decisions everywhere without threading a sink through the ops layer. + */ + public var defaultSink: TraceSink = NoopTraceSink + + private var warnedReferenceFallback: Boolean = false + /** * Select and run `matmul(a, b)`, writing into [out]. [scope] owns any adapter the selection * needs; [sink] sees the kernel run and every adapter. @@ -87,7 +129,7 @@ public object KernelDispatch { b: TensorView, out: TensorView, scope: Scope = Scope.Ambient, - sink: TraceSink = NoopTraceSink, + sink: TraceSink = defaultSink, /** * Relayout a canonical packed weight into kernel order when that is what unlocks a packed * kernel (#973/#1095). @@ -101,6 +143,9 @@ public object KernelDispatch { */ prepackWeights: Boolean = false, ) { + // Self-heal on first use: an empty table means nobody bootstrapped, and the silent + // consequence is the reference kernel for every operand pair. + ensureInstalled() val key = KernelKey.matmul(a, b) val exact = find(key) if (exact != null) { @@ -137,6 +182,21 @@ public object KernelDispatch { } // No exact kernel: adapt the operands a kernel would accept, then fall back to the reference, // which reads any format through decoding get(). + // The reference path is correct but orders of magnitude slower than a real kernel on a + // blocked weight (per-element block decode) — a process that lands here on a quantized + // weight almost certainly forgot to install a kernel pack. Say so once, loudly, even with + // no sink attached: silent fallback is how a 25 s/token regression ships unnoticed. + if (!warnedReferenceFallback && b.layout.blocked) { + warnedReferenceFallback = true + println( + "[SKaiNET] KernelDispatch: no kernel registered for matmul " + + "(activation=${a.format.encoding}, weight=${b.format.encoding}, " + + "order=${b.layout.blockOrder}); falling back to the decoding reference " + + "kernel (~1000x slower). Install a kernel pack (e.g. KernelPacks.install() " + + "+ FfmRowMajorKernelPack.install()) before the first forward. " + + "Further fallbacks are not reported." + ) + } val adaptedA = adapt(a, scope, sink, "gather") val reference = ReferenceMatmulKernel(KernelKey.matmul(adaptedA, b)) runTraced(reference, listOf(adaptedA, b), out, sink) diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.kt new file mode 100644 index 000000000..2771608d0 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.kt @@ -0,0 +1,56 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * A installable set of [ViewKernel]s for [KernelDispatch] — the view-keyed sibling of + * [KernelProvider], which serves [KernelRegistry]. + * + * Why this exists: [KernelRegistry] self-heals (an ops instance that finds it empty calls + * `KernelServiceLoader.installAll()`), but [KernelDispatch] historically did not. Every consumer + * had to remember two explicit `install()` calls before the first forward pass, and forgetting + * them is invisible: the dispatcher simply serves the decoding reference kernel, which is correct + * and roughly a thousand times slower. Making packs *discoverable* lets [KernelDispatch] populate + * itself the same way the provider registry already does. + * + * Implementations must be cheap to construct and idempotent to [install] — the dispatcher may call + * it once per process, and a pack whose native library or platform feature is unavailable should + * register nothing rather than throw. + * + * **Discovery is JVM-only**, exactly as it is for [KernelProvider]: `ServiceLoader` has no + * equivalent on Kotlin/Native, wasm or JS, so those platforms install their packs manually (see + * `installPlatformKernelPacks`). + */ +@ExperimentalMemoryApi +public interface ViewKernelPack { + + /** Stable identifier, used for logging and de-duplication (e.g. `"ffm-rowmajor"`). */ + public val name: String + + /** + * Register this pack's kernels into [KernelDispatch]. Must be safe to call more than once and + * must degrade to a no-op when the platform cannot serve it (missing native library, absent + * vector unit, …). + */ + public fun install() +} + +/** + * Discover and install every [ViewKernelPack] this platform exposes, returning the names of the + * packs that were installed. + * + * JVM: `ServiceLoader`-discovered, mirroring [KernelServiceLoader]. Everywhere else: no discovery + * mechanism exists, so this returns an empty list and the consumer installs packs explicitly. + */ +@ExperimentalMemoryApi +internal expect fun installPlatformKernelPacks(): List + +/** + * Populate [KernelRegistry] from platform-discovered [KernelProvider]s when it is still empty. + * + * Needed because [KernelPacks.install] derives its kernels from `KernelRegistry.bestAvailable()`, + * which is `null` on an empty registry — a bootstrap that runs before any ops instance exists + * would otherwise install nothing but the reference kernel. + */ +@ExperimentalMemoryApi +internal expect fun installPlatformKernelProviders(): List diff --git a/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt b/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt new file mode 100644 index 000000000..071dde841 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt @@ -0,0 +1,14 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * No `ServiceLoader` on this platform, so there is nothing to discover: packs are installed + * explicitly by the consumer (the same split [KernelProvider] documents — see + * `NativeKnKernelProvider`, which is registered by hand on Kotlin/Native). + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelPacks(): List = emptyList() + +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = emptyList() diff --git a/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.jvm.kt b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.jvm.kt new file mode 100644 index 000000000..cedf8bace --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.jvm.kt @@ -0,0 +1,26 @@ +package sk.ainet.backend.api.kernel + +import java.util.ServiceLoader +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * JVM discovery for [ViewKernelPack], mirroring [KernelServiceLoader]'s handling of + * [KernelProvider]: a backend module declares its pack in + * `META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack` and it is installed automatically + * the first time [KernelDispatch] needs kernels. + * + * A pack that throws while installing is skipped rather than allowed to break dispatch — a broken + * optional backend must not take the process down, and the reference kernel still serves every + * format correctly. + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelPacks(): List = + runCatching { + ServiceLoader.load(ViewKernelPack::class.java) + .mapNotNull { pack -> runCatching { pack.install(); pack.name }.getOrNull() } + .toList() + }.getOrElse { emptyList() } + +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = + runCatching { KernelServiceLoader.installAll() }.getOrElse { emptyList() } diff --git a/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt b/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt new file mode 100644 index 000000000..071dde841 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt @@ -0,0 +1,14 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * No `ServiceLoader` on this platform, so there is nothing to discover: packs are installed + * explicitly by the consumer (the same split [KernelProvider] documents — see + * `NativeKnKernelProvider`, which is registered by hand on Kotlin/Native). + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelPacks(): List = emptyList() + +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = emptyList() diff --git a/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt b/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt new file mode 100644 index 000000000..071dde841 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt @@ -0,0 +1,14 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * No `ServiceLoader` on this platform, so there is nothing to discover: packs are installed + * explicitly by the consumer (the same split [KernelProvider] documents — see + * `NativeKnKernelProvider`, which is registered by hand on Kotlin/Native). + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelPacks(): List = emptyList() + +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = emptyList() diff --git a/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt b/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt new file mode 100644 index 000000000..071dde841 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt @@ -0,0 +1,14 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * No `ServiceLoader` on this platform, so there is nothing to discover: packs are installed + * explicitly by the consumer (the same split [KernelProvider] documents — see + * `NativeKnKernelProvider`, which is registered by hand on Kotlin/Native). + */ +@ExperimentalMemoryApi +internal actual fun installPlatformKernelPacks(): List = emptyList() + +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = emptyList() diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 66bcbdbc3..7a638743b 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -664,14 +664,21 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val n = b.shape[1] require(k == b.shape[0]) { "Matrix multiplication shape mismatch: ${a.shape} vs ${b.shape}" } val out = FloatArray(m * n) + // Loop order i-p-j, not i-j-p. The inner statement below walks `b` and `out` + // contiguously; the j-inner-p form it replaces read `b` with stride n, touching + // a fresh cache line on nearly every multiply — for a [1536, 256] weight that is + // 1536 lines per output element. Each output still accumulates its products in + // ascending p, so the result is bit-identical, not merely close. for (i in 0 until m) { val aOff = aBase + i * k - for (j in 0 until n) { - var sum = 0f - for (p in 0 until k) { - sum += aBuf[aOff + p] * bBuf[bBase + p * n + j] + val outOff = i * n + for (p in 0 until k) { + val av = aBuf[aOff + p] + if (av == 0f) continue + val bOff = bBase + p * n + for (j in 0 until n) { + out[outOff + j] += av * bBuf[bOff + j] } - out[i * n + j] = sum } } @Suppress("UNCHECKED_CAST") @@ -937,6 +944,68 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory * than [PREPACK_CACHE_LIMIT] distinct packed weights simply converts the overflow each time, * which is exactly the old behaviour. */ + /** + * `Wᵀ` for a **dense** 2-D float weight, materialized once per weight instead of once per call. + * + * `transpose` cannot hand back a view here: the kernels want the weight input-major + * (`Fp32ViewMatmulKernel` declines anything whose `strides[0] != 1`), and no view of a + * row-major `[out, in]` buffer has that layout — only a copy does. The copy itself is a + * cache-hostile scatter over every element, so doing it per call is what actually costs: on + * Gemma 4 E2B the two dense per-layer-embedding projections are transposed 70 times per token, + * ~27M scattered element copies, and profiling attributed **84% of decode time** to those two + * matmuls while the packed projections beside them — 8x larger — ran 22x faster. + * + * This is the dense counterpart of [prepackedWeights], which has cached the packed relayout + * "once per weight instead of once per call" since #1096; the dense path simply never got the + * same treatment. Keyed on buffer identity, so it holds only for the immutable parameter a + * decode loop reuses; anything else falls through to a fresh transpose. + */ + @Suppress("UNCHECKED_CAST") + private fun transposedDenseWeight(weight: Tensor): Tensor? { + if (weight.shape.rank != 2) return null + // Key on the TensorData itself, not a FloatArray: a weight staged by the MemorySegment + // factory is not FloatArray-backed, and that is precisely the case that hurts most — + // `transpose` then falls all the way to its generic per-element fallback. + val source: Any = weight.data + transposedDenseWeights.firstOrNull { it.first === source }?.let { return it.second as Tensor } + // Materialize onto the HEAP, not through `transpose`/`dataFactory`. Two reasons: a + // MemorySegment-backed dense weight sends `transpose` to its generic per-element fallback, + // and — the bigger one — the vectorized FP32 kernels only accept heap `FloatArray` + // operands, so a segment-backed transpose is condemned to the decoding reference kernel + // afterwards. Since the cache pays for one copy anyway, it may as well land where the fast + // kernels can read it. + val transposed = heapTransposeFp32(weight) ?: transpose(weight) + if (transposedDenseWeights.size >= TRANSPOSED_DENSE_CACHE_LIMIT) transposedDenseWeights.removeAt(0) + transposedDenseWeights.add(Pair(source, transposed as Tensor<*, *>)) + return transposed + } + + /** `Wᵀ` as a heap-backed dense FP32 tensor, or null when [weight] is not dense FP32. */ + @Suppress("UNCHECKED_CAST") + private fun heapTransposeFp32(weight: Tensor): Tensor? { + if (weight.dtype != FP32::class) return null + val rows = weight.shape[0] + val cols = weight.shape[1] + val src = runCatching { weight.data.copyToFloatArray() }.getOrNull() ?: return null + if (src.size != rows * cols) return null + val out = FloatArray(src.size) + for (r in 0 until rows) { + val base = r * cols + for (c in 0 until cols) out[c * rows + r] = src[base + c] + } + return newTensor( + sk.ainet.lang.tensor.data.DenseFloatArrayTensorData(Shape(cols, rows), out) + as sk.ainet.lang.tensor.data.TensorData, + weight.dtype, + weight, + ) + } + + private val transposedDenseWeights: MutableList>> = mutableListOf() + + /** Big enough for a deep model's dense projections (Gemma 4 E2B has 70) without unbounded growth. */ + private val TRANSPOSED_DENSE_CACHE_LIMIT: Int = 256 + private val prepackedWeights: MutableList>> = mutableListOf() /** How many relayouted weights to keep; beyond this the oldest is dropped and reconverted on demand. */ @@ -963,7 +1032,9 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory @Suppress("UNCHECKED_CAST") override fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { - if (weight.shape.rank != 2 || !isHeapPackedWeight(weight.data)) return matmul(x, transpose(weight)) + if (weight.shape.rank != 2 || !isHeapPackedWeight(weight.data)) { + return matmul(x, transposedDenseWeight(weight) ?: transpose(weight)) + } // Decode the weight where it lies (#1124). The relayout below produces bytes for kernels // that address `packedData` in feed order themselves; this implementation has no such // kernel, so relayouting for it was pure harm — the result is a tensor whose shape says diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt index 201f3f93d..25b00ebd1 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorMatmulKernel.kt @@ -40,6 +40,28 @@ import sk.ainet.backend.api.kernel.Fp32MatmulKernel * scratch-pool integration is out of scope for this kernel and lives * one layer up (see `ScratchPool` SPI in `skainet-lang-core`). * + * ## Few rows take a different path + * + * Both fixed costs above are `O(n * k)` and independent of `m`: packing Bᵀ, and a horizontal + * `reduceLanes` per output cell per K-tile. They buy nothing when there are few rows to amortize + * them over, and a decode step is exactly that — `m = 1`. Measured at `k=1536, n=256` on an Apple + * M4 before [gemvRows] existed, this kernel was **slower than [ScalarMatmulKernel]**, 0.969 ms + * against 0.253. So at or below [GEMV_MAX_M] rows it accumulates by outer product instead, which + * touches B once, contiguously, and never reduces across lanes: + * + * ``` + * m tiled (before) gemvRows (after) native-ffm scalar + * 1 0.969 ms 0.81 0.058 ms 13.63 7.65 3.05 GFLOP/s + * 4 — 0.230 ms 13.70 18.50 3.07 + * 8 — 0.463 ms 13.58 21.41 3.02 + * 16 1.332 ms 9.45 (tiled) 25.53 3.06 + * 32 1.647 ms 15.28 (tiled) 27.69 3.10 + * ``` + * + * The tiled path stays in charge above that, where it is what the shape wants. Note the native FFM + * kernel wins from `m = 4` up but loses at `m = 1`, where the heap→off-heap copy a JDK 21 downcall + * requires costs more than the arithmetic; `Fp32KernelRaceBench` keeps these numbers honest. + * * Caller contract is identical to [Fp32MatmulKernel]: strides are in * floats, `out` is fully overwritten in the `m × n` block, and `k == 0` * zeros the output block. @@ -51,6 +73,12 @@ public object PanamaVectorMatmulKernel : Fp32MatmulKernel { private const val TILE_N = 8 private const val TILE_K = 128 + /** + * At or below this many rows the tiled path's O(n*k) setup outweighs its O(m*n*k) speedup, so + * [gemvRows] serves instead. Chosen from measurement, not theory — see `Fp32KernelRaceBench`. + */ + private const val GEMV_MAX_M = 8 + override fun matmul( a: FloatArray, aOffset: Int, aStride: Int, b: FloatArray, bOffset: Int, bStride: Int, @@ -70,6 +98,15 @@ public object PanamaVectorMatmulKernel : Fp32MatmulKernel { } if (k == 0) return + // Few rows: skip the tiled path entirely. Both of its fixed costs are O(n*k) — packing B + // transposed, and a horizontal reduceLanes per output cell per K-tile — so at small m they + // dwarf the O(m*n*k) arithmetic they exist to accelerate. A decode step is m=1, and there + // this kernel measured 0.969 ms against 0.253 for the scalar one it is supposed to beat. + if (m <= GEMV_MAX_M) { + gemvRows(a, aOffset, aStride, b, bOffset, bStride, out, outOffset, outStride, m, n, k) + return + } + // Pack B^T: bt[j, kk] = b[kk, j]. Row stride in bt is k. val bt = FloatArray(n * k) for (kk in 0 until k) { @@ -103,6 +140,45 @@ public object PanamaVectorMatmulKernel : Fp32MatmulKernel { } } + /** + * Outer-product accumulation over `out` rows: `out[i, :] += a[i, p] * b[p, :]`. + * + * Streams `b` and `out` contiguously along `n` and never transposes or packs anything, so the + * whole call costs one pass over B. There is no horizontal reduction — each lane owns one + * output column for the entire `k` loop — which is the other thing the tiled path pays per + * cell. Accumulation for a given output stays in ascending `p`, matching the scalar kernel's + * order rather than the tiled path's split-by-K-tile order. + */ + private fun gemvRows( + a: FloatArray, aOffset: Int, aStride: Int, + b: FloatArray, bOffset: Int, bStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + val step = species.length() + val bound = species.loopBound(n) + for (i in 0 until m) { + val aBase = aOffset + i * aStride + val outRow = outOffset + i * outStride + for (p in 0 until k) { + val av = a[aBase + p] + if (av == 0f) continue + val vav = FloatVector.broadcast(species, av) + val bRow = bOffset + p * bStride + var j = 0 + while (j < bound) { + val vo = FloatVector.fromArray(species, out, outRow + j) + vav.fma(FloatVector.fromArray(species, b, bRow + j), vo).intoArray(out, outRow + j) + j += step + } + while (j < n) { + out[outRow + j] += av * b[bRow + j] + j++ + } + } + } + } + /** * Recursive (m, n) tile dispatch. Picks the largest microkernel * shape `(RM, RN)` that fits the residual `(m1-m0, n1-n0)`, calls it diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 36bffd447..a4b3fcdbf 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -198,9 +198,11 @@ internal class DefaultCpuOpsJvm( // SegmentStorage by design). On this tier the bridge is trivial: bulk-copy the // activation to the heap once per call — decode-step activations are k floats, not // weights — and hand the common path heap views. - if (weight.data is sk.ainet.lang.tensor.storage.PackedBlockStorage) { - segmentActivationToHeap(x)?.let { return super.matmulWeightTransposed(it, weight) } - } + // Dense weights need this bridge just as much as packed ones. The vectorized FP32 + // kernels take heap `FloatArray` operands only, so a segment-backed activation sends + // the whole projection to the decoding reference kernel — and a decode-step activation + // is k floats, so the copy is trivial beside the matmul it unlocks. + segmentActivationToHeap(x)?.let { return super.matmulWeightTransposed(it, weight) } return super.matmulWeightTransposed(x, weight) } // Already in feed order — the loader produced it that way (#1120). Nothing to permute and @@ -254,6 +256,14 @@ internal class DefaultCpuOpsJvm( private val prepackedWeightsJvm: MutableList>> = mutableListOf() private val PREPACK_CACHE_LIMIT_JVM: Int = 64 + /** + * Below this many multiply-accumulates a dense FP32 matmul is faster done directly than handed + * to the tile-blocked SPI kernel, whose per-call setup dominates at these sizes. Sized to cover + * decode-step projections (m=1) while leaving prefill batches and the big vocab matmuls to the + * kernel — at m=32 the kernel is already ~16 GFLOP/s and pulling away. + */ + private val SMALL_FP32_MATMUL_WORK: Long = 4_000_000L + /** The heap packed types whose JVM kernels read input-block-major bytes. */ private fun isHeapPackedWeightForJvm(data: sk.ainet.lang.tensor.data.TensorData<*, *>): Boolean = data is sk.ainet.lang.tensor.data.Q4_KTensorData || data is sk.ainet.lang.tensor.data.Q5_KTensorData || @@ -1097,6 +1107,33 @@ internal class DefaultCpuOpsJvm( } } + // Small shapes are overhead-bound, not throughput-bound. The tile-blocked SPI kernel below + // costs ~1ms per call whether it multiplies 1 row or 8 (Fp32GemvShapeBench: m=1 1.00ms, + // m=8 1.08ms — 8x the arithmetic for 8% more time), so at decode sizes essentially all of + // that is fixed cost. A decode step is m=1 by construction, and a model whose checkpoint + // ships dense FP32 tensors does this per layer: Gemma 4 E2B has 70 such projections per + // token and spent 73% of decode here. Do the small ones directly instead — contiguous in + // `b` and `out`, one pass, no setup — and leave the tiled kernel the large shapes it wins. + if (work <= SMALL_FP32_MATMUL_WORK) { + val aBuf = aWin.arr + val aBase = aWin.off + val bBuf = bWin.arr + val bBase = bWin.off + for (i in 0 until m) { + val aOff = aBase + i * k + val outOff = i * n + for (p in 0 until k) { + val av = aBuf[aOff + p] + if (av == 0f) continue + val bOff = bBase + p * n + for (j in 0 until n) { + outBuffer[outOff + j] += av * bBuf[bOff + j] + } + } + } + return floatResult(Shape(m, n), a.dtype, outBuffer) + } + // Route through the kernel SPI — the registered provider // (Panama on JDK 21+, scalar otherwise) is tile-blocked and // handles small + large inputs in one path, so the previous diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp32GemvShapeBench.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp32GemvShapeBench.kt new file mode 100644 index 000000000..b7a459f83 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp32GemvShapeBench.kt @@ -0,0 +1,83 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.matmul +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.time.measureTime + +/** + * How the dense FP32 matmul performs as a function of `m` — i.e. GEMV (`m = 1`, one decode step) + * versus GEMM (`m > 1`, a prefill batch). + * + * Why this exists: decode is entirely `m = 1`, and the FP32 SPI kernel is tile-blocked, a shape + * GEMM kernels are usually poorest at. Profiling Gemma 4 E2B put **73% of decode** in two dense + * FP32 projections per layer — 0.39M MACs each, yet ~4x more expensive per call than a packed Q4_K + * projection 8x their size. Those weights are dense because the checkpoint ships them unquantized + * (70 of its 2-D tensors are F32: the per-layer-embedding gate and projection), so this shape is + * not exotic — any model with higher-precision tensors lands here. + * + * Not an assertion, a measurement: run it and read the numbers. + * `./gradlew :skainet-backends:skainet-backend-cpu:jvmTest --tests "*Fp32GemvShapeBench*" -i` + */ +class Fp32GemvShapeBench { + + @Test + fun gemv_versus_gemm_throughput() { + if (System.getenv("SKAINET_BENCH") != "1") { + println("[skip] set SKAINET_BENCH=1 to run the FP32 shape benchmark"); return + } + val ctx = DirectCpuExecutionContext() + val k = 1536 // Gemma 4 E2B hidden size + val n = 256 // per-layer-embedding width + val iterations = 200 + + for (m in intArrayOf(1, 2, 4, 8, 16, 32)) { + val a = ctx.fromFloatArray( + Shape(m, k), FP32::class, FloatArray(m * k) { (it % 17) * 0.03f }, + ) + val b = ctx.fromFloatArray( + Shape(k, n), FP32::class, FloatArray(k * n) { (it % 13) * 0.02f }, + ) + repeat(20) { a.matmul(b) } // warm up JIT + val elapsed = measureTime { repeat(iterations) { a.matmul(b) } } + val macs = m.toLong() * k * n * iterations + val gflops = 2.0 * macs / elapsed.inWholeNanoseconds + println( + "BENCH m=%-3d %7.3f ms/call %6.2f GFLOP/s %6.3f ms per output row".format( + m, + elapsed.inWholeMicroseconds / 1000.0 / iterations, + gflops, + elapsed.inWholeMicroseconds / 1000.0 / iterations / m, + ) + ) + } + println("BENCH reference: the packed Q4_K path measures ~29 GFLOP/s on this machine") + + // Same shapes with BOTH operands off-heap, which is what a read-only weight can simply be: + // `chooseMatmul` then takes its MemorySegment branch and the kernel reads the weight in + // place instead of copying it off-heap on every call. + val msCtx = DirectCpuExecutionContext( + tensorDataFactory = sk.ainet.lang.tensor.data.MemorySegmentTensorDataFactory(), + ) + for (m in intArrayOf(1, 8, 32)) { + val a = msCtx.fromFloatArray( + Shape(m, k), FP32::class, FloatArray(m * k) { (it % 17) * 0.03f }, + ) + val b = msCtx.fromFloatArray( + Shape(k, n), FP32::class, FloatArray(k * n) { (it % 13) * 0.02f }, + ) + println("BENCH memseg operands: a=${a.data::class.simpleName} b=${b.data::class.simpleName}") + repeat(20) { a.matmul(b) } + val elapsed = measureTime { repeat(iterations) { a.matmul(b) } } + val macs = m.toLong() * k * n * iterations + println( + "BENCH memseg m=%-3d %7.3f ms/call %6.2f GFLOP/s".format( + m, elapsed.inWholeMicroseconds / 1000.0 / iterations, + 2.0 * macs / elapsed.inWholeNanoseconds, + ) + ) + } + } +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniMappedKernelPackFactory.kt b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniMappedKernelPackFactory.kt new file mode 100644 index 000000000..8bb5e2a29 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniMappedKernelPackFactory.kt @@ -0,0 +1,20 @@ +package sk.ainet.exec.kernel.jni + +import sk.ainet.backend.api.kernel.ViewKernelPack +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * `ServiceLoader`-friendly wrapper around [JniMappedKernelPack], the Android counterpart of + * `FfmRowMajorKernelPackFactory`: it lets `KernelDispatch.ensureInstalled()` discover the JNI + * row-major kernels so an Android consumer gets zero-copy mapped weights without an explicit + * bootstrap call. + * + * Listed in `META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack`. Note that Android + * packaging must preserve `META-INF/services` entries for discovery to work; a consumer whose + * build strips them can still call [JniMappedKernelPack.install] directly. + */ +@OptIn(ExperimentalMemoryApi::class) +public class JniMappedKernelPackFactory : ViewKernelPack { + override val name: String get() = "jni-rowmajor" + override fun install(): Unit = JniMappedKernelPack.install() +} diff --git a/skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack b/skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack new file mode 100644 index 000000000..f396ddd25 --- /dev/null +++ b/skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack @@ -0,0 +1 @@ +sk.ainet.exec.kernel.jni.JniMappedKernelPackFactory diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernelPackFactory.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernelPackFactory.kt new file mode 100644 index 000000000..0d6555ec0 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernelPackFactory.kt @@ -0,0 +1,20 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.ViewKernelPack +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * `ServiceLoader`-friendly wrapper around [FfmRowMajorKernelPack] — the same shape + * [NativeKernelProviderFactory] gives [NativeKernelProvider], because `ServiceLoader` needs a + * public no-arg constructor and a Kotlin `object` does not expose one. + * + * Listed in `META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack` so + * `KernelDispatch.ensureInstalled()` finds the FFM row-major kernels without the consumer calling + * anything. [FfmRowMajorKernelPack.install] is already a no-op when the native library is missing, + * so discovery on a machine without it costs a lookup and registers nothing. + */ +@OptIn(ExperimentalMemoryApi::class) +public class FfmRowMajorKernelPackFactory : ViewKernelPack { + override val name: String get() = "ffm-rowmajor" + override fun install(): Unit = FfmRowMajorKernelPack.install() +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt index 509a00e9c..2cf5b66b9 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp32MatmulKernel.kt @@ -34,6 +34,44 @@ import sk.ainet.backend.api.kernel.Fp32MatmulKernel * future work could add parallelChunks-style row blocking and B-tile * packing, but the scalar C path already lands well within the SPI * contract on host-arch CPUs. + * + * ## Per-call cost, and why callers avoid this kernel at small sizes + * + * The SPI hands this kernel heap `FloatArray`s, and a downcall cannot address heap memory on + * JDK 21, so both operands are copied off-heap on every call. That copy is proportional to the + * *weight*, not to the work: at `k=1536, n=256` it is 1.5 MB whatever `m` is. Measured on an + * Apple M4 (`Fp32GemvShapeBench`): + * + * ``` + * m= 1 1.002 ms/call 0.78 GFLOP/s m= 8 1.075 ms/call 5.85 GFLOP/s + * m= 4 0.992 ms/call 3.17 GFLOP/s m=32 1.600 ms/call 15.73 GFLOP/s + * ``` + * + * Eight times the arithmetic for 8% more time — below `m ~ 32` this is a fixed cost, so a decode + * step (`m = 1`) is essentially all copy. Two things were tried and measured as no-ops, so do not + * reach for them again: reusing the [Arena] and its segments across calls (1.025 ms vs 1.002 ms at + * m=1 — the allocation was never the cost), and reordering the C loops (the C kernel is already + * i-p-j). `DefaultCpuOpsJvm` therefore serves small shapes directly and leaves this kernel the + * large ones it wins. + * + * Keeping the weight off-heap so it needs no copy is the obvious escape — a weight is read-only + * for the life of the process — and it does not pay today, because the segment kernel that then + * serves it (`JvmVectorKernels.matmulFloatBlockedMemSeg`) is slower than the heap one by more than + * the copy costs. Measured on the same shapes, both operands `MemorySegmentTensorData`: + * + * ``` + * m= 1 0.943 ms/call 0.83 GFLOP/s (heap path: 0.122 ms, 6.44) + * m= 8 1.363 ms/call 4.62 GFLOP/s (heap path: 0.770 ms, 8.17) + * m=32 2.785 ms/call 9.04 GFLOP/s (heap path: 1.618 ms, 15.55) + * ``` + * + * So there are two independent gaps, and residency is not the lever: this kernel cannot see heap + * memory, and the kernel that can see off-heap memory is slow. Closing either one is worth real + * throughput — the packed Q4_K path next door reaches ~29 GFLOP/s on the same machine. + * + * The copy disappears on **JDK 22+**, where `Linker.Option.critical(true)` lets a downcall read + * heap segments directly; when the toolchain moves, pass `MemorySegment.ofArray(...)` through a + * critical handle and the small-shape threshold in `DefaultCpuOpsJvm` can be revisited. */ internal object NativeFp32MatmulKernel : Fp32MatmulKernel { diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack new file mode 100644 index 000000000..29450fd1a --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack @@ -0,0 +1 @@ +sk.ainet.exec.kernel.FfmRowMajorKernelPackFactory diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Fp32KernelRaceBench.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Fp32KernelRaceBench.kt new file mode 100644 index 000000000..301bb438f --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Fp32KernelRaceBench.kt @@ -0,0 +1,85 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.KernelProvider +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.backend.api.kernel.KernelServiceLoader +import kotlin.test.Test +import kotlin.time.measureTime + +/** + * Races every registered FP32 matmul kernel against the same decode-shaped problem, so "the native + * one is slower" becomes a number per provider rather than an inference from end-to-end timings. + * + * Context: at `m=1, k=1536, n=256` the SPI kernel costs ~1.00 ms/call. The heap→off-heap copy the + * native kernel performs was the obvious suspect and is not: measured separately at 0.029 ms for + * the same 1.5 MB (54 GB/s, `HeapToSegmentCopyBench`). That leaves the kernels' own compute. + * + * Also times the fp16 kernel where one is registered: half precision is plausibly *faster* rather + * than merely smaller on AArch64, which doubles FP16 FLOPs under FEAT_FP16, and these particular + * weights (a per-layer-embedding gate/projection) are a side channel where the precision is + * affordable. + */ +class Fp32KernelRaceBench { + + @Test + fun race_registered_fp32_kernels() { + if (System.getenv("SKAINET_BENCH") != "1") { + println("[skip] set SKAINET_BENCH=1 to run the kernel race"); return + } + if (KernelRegistry.providers().isEmpty()) KernelServiceLoader.installAll() + + val k = 1536 + val n = 256 + val iterations = 300 + for (m in intArrayOf(1, 4, 8, 16, 32)) raceAt(m, k, n, iterations) + } + + private fun raceAt(m: Int, k: Int, n: Int, iterations: Int) { + val a = FloatArray(m * k) { (it % 17) * 0.03f } + val b = FloatArray(k * n) { (it % 13) * 0.02f } + val out = FloatArray(m * n) + val macs = m.toLong() * k * n * iterations + + println("RACE ---- m=$m k=$k n=$n ----") + + for (p: KernelProvider in KernelRegistry.providers()) { + if (!p.isAvailable()) continue + val kernel = p.matmulFp32() ?: run { println("RACE %-16s fp32: ".format(p.name)); null } ?: continue + repeat(30) { kernel.matmul(a, 0, k, b, 0, n, out, 0, n, m, n, k) } + val elapsed = measureTime { + repeat(iterations) { kernel.matmul(a, 0, k, b, 0, n, out, 0, n, m, n, k) } + } + println( + "RACE %-16s fp32 %7.3f ms/call %6.2f GFLOP/s [%s]".format( + p.name, elapsed.inWholeMicroseconds / 1000.0 / iterations, + 2.0 * macs / elapsed.inWholeNanoseconds, kernel::class.simpleName, + ) + ) + } + + // The plain Kotlin loop DefaultCpuOpsJvm now uses for small shapes, for comparison. + repeat(30) { directLoop(a, b, out, m, n, k) } + val direct = measureTime { repeat(iterations) { directLoop(a, b, out, m, n, k) } } + println( + "RACE %-16s fp32 %7.3f ms/call %6.2f GFLOP/s".format( + "kotlin-direct", direct.inWholeMicroseconds / 1000.0 / iterations, + 2.0 * macs / direct.inWholeNanoseconds, + ) + ) + + } + + private fun directLoop(a: FloatArray, b: FloatArray, out: FloatArray, m: Int, n: Int, k: Int) { + java.util.Arrays.fill(out, 0f) + for (i in 0 until m) { + val aOff = i * k + val outOff = i * n + for (p in 0 until k) { + val av = a[aOff + p] + if (av == 0f) continue + val bOff = p * n + for (j in 0 until n) out[outOff + j] += av * b[bOff + j] + } + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/HeapToSegmentCopyBench.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/HeapToSegmentCopyBench.kt new file mode 100644 index 000000000..b50e3e959 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/HeapToSegmentCopyBench.kt @@ -0,0 +1,72 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.nio.ByteOrder +import kotlin.test.Test +import kotlin.time.measureTime + +/** + * How fast is heap `FloatArray` → off-heap `MemorySegment`, really? + * + * `NativeFp32MatmulKernel` copies both operands off-heap on every call because a JDK 21 downcall + * cannot address heap memory. Decomposing its measured cost (1.002 ms at m=1, 1.600 ms at m=32, + * k=1536 n=256) as `total = copy + compute·m` puts compute at ~19 µs — the C kernel is doing + * ~41 GFLOP/s — and the copy at ~0.98 ms for 1.5 MB, i.e. ~1.5 GB/s. That is an order of magnitude + * under this machine's memory bandwidth, so the question is whether the copy is intrinsified at all + * or is quietly running element-wise. + * + * Compares the spellings available on JDK 21 for the same 1.5 MB. + */ +class HeapToSegmentCopyBench { + + @Test + fun heap_to_offheap_copy_throughput() { + if (System.getenv("SKAINET_BENCH") != "1") { + println("[skip] set SKAINET_BENCH=1 to run the copy benchmark"); return + } + val floats = 1536 * 256 // the Gemma 4 per-layer-embedding weight + val bytes = floats.toLong() * Float.SIZE_BYTES + val src = FloatArray(floats) { it * 0.001f } + val iterations = 500 + + fun report(label: String, elapsedMs: Double) { + val gbPerSec = bytes.toDouble() * iterations / (elapsedMs / 1000.0) / 1e9 + println("COPY %-42s %7.3f ms/copy %6.2f GB/s".format(label, elapsedMs / iterations, gbPerSec)) + } + + Arena.ofConfined().use { arena -> + val dst = arena.allocate(bytes, ValueLayout.JAVA_FLOAT.byteAlignment()) + + // 1. What the kernel does today. + repeat(50) { MemorySegment.copy(src, 0, dst, ValueLayout.JAVA_FLOAT, 0L, floats) } + report("MemorySegment.copy(JAVA_FLOAT)", measureTime { + repeat(iterations) { MemorySegment.copy(src, 0, dst, ValueLayout.JAVA_FLOAT, 0L, floats) } + }.inWholeMicroseconds / 1000.0) + + // 2. Same, with the layout's byte order pinned to native. An unaligned or non-native + // layout is the usual reason this drops off its intrinsic. + val nativeLayout = ValueLayout.JAVA_FLOAT.withOrder(ByteOrder.nativeOrder()) + repeat(50) { MemorySegment.copy(src, 0, dst, nativeLayout, 0L, floats) } + report("MemorySegment.copy(JAVA_FLOAT native order)", measureTime { + repeat(iterations) { MemorySegment.copy(src, 0, dst, nativeLayout, 0L, floats) } + }.inWholeMicroseconds / 1000.0) + + // 3. Segment-to-segment bulk copy, wrapping the heap array as a segment. This is the + // memcpy-shaped spelling and does not go through a ValueLayout at all. + val srcSeg = MemorySegment.ofArray(src) + repeat(50) { dst.copyFrom(srcSeg) } + report("dst.copyFrom(MemorySegment.ofArray(src))", measureTime { + repeat(iterations) { dst.copyFrom(srcSeg) } + }.inWholeMicroseconds / 1000.0) + + // 4. Floor: heap-to-heap arraycopy, for scale. + val heapDst = FloatArray(floats) + repeat(50) { System.arraycopy(src, 0, heapDst, 0, floats) } + report("System.arraycopy (heap->heap, for scale)", measureTime { + repeat(iterations) { System.arraycopy(src, 0, heapDst, 0, floats) } + }.inWholeMicroseconds / 1000.0) + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelDispatchSelfHealTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelDispatchSelfHealTest.kt new file mode 100644 index 000000000..0753351c7 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelDispatchSelfHealTest.kt @@ -0,0 +1,64 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.lang.memory.ExperimentalMemoryApi +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * The dispatcher must populate itself on first use, with no bootstrap call from the consumer. + * + * Before this, `KernelDispatch` was the only registry that did not self-heal: `KernelRegistry` + * installs providers lazily via `DefaultCpuOpsJvm.ensureKernelProviders()`, but every dispatch + * consumer had to remember `KernelPacks.install()` + `FfmRowMajorKernelPack.install()` before its + * first forward pass. Forgetting it is invisible — the decoding reference kernel is correct, just + * ~1000x slower — and it was forgotten by application entry points and diagnostic harnesses alike. + */ +@OptIn(ExperimentalMemoryApi::class) +class KernelDispatchSelfHealTest { + + @Test + fun cold_dispatch_installs_providers_and_discovered_packs() { + KernelDispatch.clearForTesting() + KernelRegistry.clearForTesting() + + // No bootstrap of any kind — exactly what a forgetful consumer does. + KernelDispatch.ensureInstalled() + + val names = KernelDispatch.kernels().map { it.name } + assertTrue(names.isNotEmpty(), "self-heal must register kernels") + assertTrue( + KernelRegistry.providers().isNotEmpty(), + "providers must be discovered first, since KernelPacks.install() derives from them", + ) + assertTrue( + names.any { it.endsWith("-fp32") }, + "provider-derived dense FP32 view kernels expected; got $names", + ) + assertTrue( + names.any { it.startsWith("ffm-rowmajor-") }, + "ServiceLoader-discovered ViewKernelPack (FFM row-major) expected; got $names", + ) + assertTrue( + KernelDispatch.mappedServableEncodings().isNotEmpty(), + "row-major pack should make mapped K-quant weights servable zero-copy", + ) + println("SELFHEAL n=${names.size} providers=${KernelRegistry.availableNames()} kernels=${names.sorted()}") + } + + @Test + fun explicit_registration_suppresses_auto_install() { + KernelDispatch.clearForTesting() + KernelRegistry.clearForTesting() + // A consumer that wires its own kernels keeps full control: auto-install must not run + // behind its back and silently add tiers it deliberately left out. + sk.ainet.backend.api.kernel.KernelPacks.installReference() + val afterExplicit = KernelDispatch.kernels().map { it.name } + KernelDispatch.ensureInstalled() + assertTrue( + KernelDispatch.kernels().map { it.name } == afterExplicit, + "ensureInstalled() must be a no-op once the table is non-empty", + ) + } +} diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/SpecialTokenSplitter.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/SpecialTokenSplitter.kt index d1cd5d7d4..90c5421fb 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/SpecialTokenSplitter.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/SpecialTokenSplitter.kt @@ -77,6 +77,19 @@ public class SpecialTokenSplitter( return IntArray(out.size) { out[it] } } + /** + * Single-id decode must delegate to the base's own [Tokenizer.decodeToken], + * not to [decode] — the interface default would route through this + * decorator's batch path, which calls `base.decode(ids)` and thereby + * re-enables leading-space stripping that bases like + * [SentencePieceTokenizer] deliberately disable for per-token streaming. + * Without this override every SentencePiece GGUF with specials (all + * Gemma-family chat models) loses word-boundary spaces when decoded + * token-by-token: "the process" streams as "theprocess". + */ + override fun decodeToken(id: Int): String = + specialIdToString[id] ?: base.decodeToken(id) + override fun decode(ids: IntArray): String { if (ids.isEmpty()) return "" if (specialTokens.isEmpty()) return base.decode(ids) diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt index 6dbb8a5f4..535100872 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/tokenizer/TokenizerFactory.kt @@ -46,7 +46,10 @@ public object TokenizerFactory { ) return when (model) { "gpt2", "bpe" -> QwenByteLevelBpeTokenizer.fromGgufFields(fields) - "llama", "sentencepiece" -> wrapSentencePieceWithSpecialsFromGguf( + // "gemma4": Gemma 4 GGUFs declare their own model string but carry a + // standard SentencePiece vocab with CONTROL/USER_DEFINED specials + // (<|turn>, , tool markers) — same shape as "llama". + "llama", "sentencepiece", "gemma4" -> wrapSentencePieceWithSpecialsFromGguf( base = SentencePieceTokenizer.fromGgufFields(fields), fields = fields, ) @@ -124,8 +127,11 @@ public object TokenizerFactory { ): Tokenizer { val tokens = (fields["tokenizer.ggml.tokens"] as? List<*>) ?.filterIsInstance().orEmpty() + // toIntFlexible, not `as? Number`: GGUF UINT32/INT32 arrays surface as + // kotlin.UInt (a value class, not Number) — the plain cast silently + // yields an empty list, dropping every special token on such files. val tokenTypes = (fields["tokenizer.ggml.token_type"] as? List<*>) - ?.mapNotNull { (it as? Number)?.toInt() }.orEmpty() + ?.mapNotNull { it.toIntFlexible() }.orEmpty() if (tokens.isEmpty() || tokenTypes.isEmpty()) return base val specials = HashMap() diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/registry/ModelArchitecture.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/registry/ModelArchitecture.kt index bb8bdb823..aad679577 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/registry/ModelArchitecture.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/registry/ModelArchitecture.kt @@ -7,6 +7,7 @@ import sk.ainet.io.model.ModelArchitecture */ private val ggufIdMap: Map = mapOf( "llama" to ModelArchitecture.LLAMA, + "gemma4" to ModelArchitecture.GEMMA, "gemma3n" to ModelArchitecture.GEMMA, "gemma3" to ModelArchitecture.GEMMA, "gemma" to ModelArchitecture.GEMMA,