From 220c77dedffea512f3eeba1b8096ce542f579bcf Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 30 Aug 2026 22:52:58 +0200 Subject: [PATCH 1/8] fix(io,backend): gemma4 tokenizer/arch registry, streaming decodeToken, observable dispatch fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four fixes surfaced by the Gemma 4 E2B GGUF debugging session (SKaiNET-transformers#325/#330 arc): - TokenizerFactory: accept 'gemma4' as tokenizer.ggml.model (SentencePiece + specials wrap) — the allowlist threw UnsupportedTokenizerException on real Gemma 4 GGUFs even though the pipeline handles them. - TokenizerFactory: read tokenizer.ggml.token_type via toIntFlexible — GGUF UINT32 arrays surface as kotlin.UInt (not Number), so the old 'as? Number' filter silently dropped every special token on such files. - SpecialTokenSplitter: override decodeToken to delegate to the base's decodeToken — the interface default routed through the batch decode path, re-enabling leading-space stripping and losing word-boundary spaces in per-token streaming for every SentencePiece GGUF with specials. - KernelDispatch: settable defaultSink + a one-time loud warning when the decoding reference matmul serves a blocked weight — the ~1000x fallback used to be invisible because production call sites pass no TraceSink. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../backend/api/kernel/KernelDispatch.kt | 27 ++++++++++++++++++- .../io/tokenizer/SpecialTokenSplitter.kt | 13 +++++++++ .../sk/ainet/io/tokenizer/TokenizerFactory.kt | 10 +++++-- .../io/gguf/registry/ModelArchitecture.kt | 1 + 4 files changed, 48 insertions(+), 3 deletions(-) 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..065cd2b42 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 @@ -76,6 +76,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 +97,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). @@ -137,6 +147,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-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, From 3734006f84abad611937daeacb9a744b4be8bed9 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 08:10:50 +0200 Subject: [PATCH 2/8] feat(kernel): KernelDispatch self-heals via a discoverable ViewKernelPack SPI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KernelDispatch was the only registry that did not populate itself. KernelRegistry self-heals (DefaultCpuOpsJvm.ensureKernelProviders installs providers lazily on first use), 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 for every format, just about a thousand times slower. It was forgotten repeatedly in practice — by an application entry point (SKaiNET-transformers' kgemma CLI) and by a diagnostic harness, where it turned a ten-forward-pass run into 23 minutes of CPU. - New ViewKernelPack SPI in skainet-backend-api: an installable set of ViewKernels, the view-keyed sibling of KernelProvider. Discovery is ServiceLoader-based on JVM and Android and manual elsewhere, the same split KernelProvider already documents (expect/actual across all six source sets). - KernelDispatch.ensureInstalled(), called at the top of matmul(): when nothing is registered, discover providers, then KernelPacks.install(), then every discovered pack. Providers must come first — KernelPacks.install() derives its kernels from KernelRegistry.bestAvailable(), which is null on an empty registry, so the reverse order silently installs only the reference kernel (measured: 8 of 17 kernels). - FfmRowMajorKernelPackFactory (native-cpu) and JniMappedKernelPackFactory (jni-cpu) plus their META-INF/services entries, so the row-major kernels that serve mapped GGUF weights zero-copy are found without a consumer call. - Explicit registration still wins: ensureInstalled() is a no-op once the table is non-empty, so a curated kernel set is never silently widened. Verified end-to-end: with the downstream bootstrap removed entirely, a Gemma 4 E2B Q4_K_M parity run loads and does 10 forward passes in 8.7s and produces the same output as with an explicit bootstrap. Docs: new explanation/kernel-selection.adoc covering both registries, the SPI and its registration/discovery rules, the selection algorithm (including the adapter and reference fallbacks), why a kernel declines, and how to diagnose which kernel ran. Cross-linked from explanation/eager-execution.adoc, which covers the provider tier only. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- docs/modules/ROOT/nav.adoc | 1 + .../pages/explanation/eager-execution.adoc | 5 + .../pages/explanation/kernel-selection.adoc | 214 ++++++++++++++++++ .../api/kernel/ViewKernelPack.android.kt | 21 ++ .../backend/api/kernel/KernelDispatch.kt | 37 ++- .../backend/api/kernel/ViewKernelPack.kt | 56 +++++ .../api/kernel/ViewKernelPack.other.kt | 14 ++ .../backend/api/kernel/ViewKernelPack.jvm.kt | 26 +++ .../api/kernel/ViewKernelPack.other.kt | 14 ++ .../api/kernel/ViewKernelPack.other.kt | 14 ++ .../api/kernel/ViewKernelPack.other.kt | 14 ++ .../kernel/jni/JniMappedKernelPackFactory.kt | 20 ++ ...sk.ainet.backend.api.kernel.ViewKernelPack | 1 + .../kernel/FfmRowMajorKernelPackFactory.kt | 20 ++ ...sk.ainet.backend.api.kernel.ViewKernelPack | 1 + .../exec/kernel/KernelDispatchSelfHealTest.kt | 64 ++++++ 16 files changed, 521 insertions(+), 1 deletion(-) create mode 100644 docs/modules/ROOT/pages/explanation/kernel-selection.adoc create mode 100644 skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.android.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.kt create mode 100644 skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt create mode 100644 skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.jvm.kt create mode 100644 skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt create mode 100644 skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt create mode 100644 skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.other.kt create mode 100644 skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniMappedKernelPackFactory.kt create mode 100644 skainet-backends/skainet-backend-jni-cpu/src/main/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/FfmRowMajorKernelPackFactory.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmMain/resources/META-INF/services/sk.ainet.backend.api.kernel.ViewKernelPack create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/KernelDispatchSelfHealTest.kt 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..2f60a006c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.android.kt @@ -0,0 +1,21 @@ +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() } + +@ExperimentalMemoryApi +internal actual fun installPlatformKernelProviders(): List = + runCatching { KernelServiceLoader.installAll() }.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 065cd2b42..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` @@ -111,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) { 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-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/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/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", + ) + } +} From 8d4456f1ff357a8d11d8ee07ce8894ce15691241 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 11:02:39 +0200 Subject: [PATCH 3/8] perf(cpu): cache the dense weight transpose, and land it where the fast kernels can read it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit matmulWeightTransposed materialized Wt on EVERY call for a dense weight, while the packed path has cached its relayout 'once per weight instead of once per call' since #1096. transpose() cannot return a view here — the FP32 kernels decline any weight whose strides[0] != 1, and no view of a row-major [out, in] buffer has that layout — so each call paid a full cache-hostile scatter over every element. Three changes, all in that one asymmetry: - cache Wt per weight, keyed on TensorData identity (not a FloatArray: a weight staged by the MemorySegment factory is not FloatArray-backed, and that is the case that hurts most — transpose then falls to its generic per-element fallback); - materialize the cached copy on the HEAP rather than through dataFactory, because the vectorized FP32 kernels take heap FloatArray operands only, so a segment-backed transpose is condemned to the reference kernel afterwards; - widen the JVM activation->heap bridge from packed weights to all weights, for the same reason on the other operand. Measured on Gemma 4 E2B Q4_K_M (its per-layer-embedding projections are the only dense-FP32 weights in an otherwise Q4_K model, transposed 70x per token): decode 443 -> 272 ms/token (2.26 -> 3.68 tok/s) prefill 432 -> 231 ms/token (2.31 -> 4.32 tok/s) Correctness unchanged: the Gemma 4 golden-token gate, its GGUF smoke test, FunctionGemma's tokenizer/template/e2e gates and the engine backend suites all pass. Note the answers move in the last ulp — a vectorized kernel accumulates in a different order than the scalar loop it replaces. Still open: even heap-on-heap, a dense FP32 projection runs ~0.3 GFLOP/s against ~29 for the packed Q4_K path beside it, so those two matmuls remain 73% of decode. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 66 ++++++++++++++++++- .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 8 ++- 2 files changed, 70 insertions(+), 4 deletions(-) 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..f2321d844 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 @@ -937,6 +937,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") + protected 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 +1025,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/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 36bffd447..2b81e6eea 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 From 558a31a6ea88db96aabd1c3ef3f4090a8f94e3ea Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 11:16:30 +0200 Subject: [PATCH 4/8] =?UTF-8?q?perf(cpu):=20serve=20small=20dense=20FP32?= =?UTF-8?q?=20matmuls=20directly=20=E2=80=94=20the=20tiled=20kernel=20is?= =?UTF-8?q?=20overhead-bound=20there?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FP32 SPI kernel is tile-blocked, and at decode sizes its per-call setup is essentially the whole cost. Measured (Fp32GemvShapeBench, k=1536 n=256): m= 1 1.002 ms/call 0.78 GFLOP/s m= 8 1.075 ms/call 5.85 GFLOP/s m= 2 1.028 ms/call 1.53 GFLOP/s m=16 1.260 ms/call 9.98 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 ~1ms overhead, not a throughput curve. Decode is m=1 by construction, so a model whose checkpoint ships dense FP32 tensors pays it per layer per token — Gemma 4 E2B has 70 such projections (its per-layer-embedding gate and proj are F32 in the GGUF while everything around them is Q4_K) and spent 73% of decode there. Small shapes now go through a direct i-p-j loop instead: contiguous in b and out, one pass, no setup. The threshold (4M MACs) covers decode-step projections and leaves prefill batches and the vocab matmul to the tiled kernel, which is already ~16 GFLOP/s at m=32 and pulling away. m=1 after: 0.125 ms/call, 6.30 GFLOP/s — 8x End to end on Gemma 4 E2B Q4_K_M, cumulative with the transpose-cache commit: decode 443 -> 189 ms/token (2.26 -> 5.29 tok/s) 2.3x prefill 432 -> 186 ms/token (2.31 -> 5.37 tok/s) 2.3x the two PLE projections: 1518 -> 49 ms over 8 tokens 31x and the decode profile is balanced again: lm_head 29%, attention 28%, RoPE 14%, PLE 15%, where PLE alone used to be 73%. Gates green: Gemma 4 golden-token and GGUF smoke, FunctionGemma tokenizer/template /e2e, tokenizer parity, and the engine backend suites. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 17 ++++-- .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 35 +++++++++++ .../exec/tensor/ops/Fp32GemvShapeBench.kt | 58 +++++++++++++++++++ 3 files changed, 105 insertions(+), 5 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp32GemvShapeBench.kt 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 f2321d844..e3e55eeb6 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") 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 2b81e6eea..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 @@ -256,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 || @@ -1099,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..8bdedc4c0 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp32GemvShapeBench.kt @@ -0,0 +1,58 @@ +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") + } +} From 019c664068ad0bba559e29475b48bdc91535d5b3 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 11:29:44 +0200 Subject: [PATCH 5/8] docs(kernel): record why small dense FP32 matmuls bypass the native kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three paths measured on an Apple M4 at k=1536 n=256 (Fp32GemvShapeBench), so the next person does not re-derive them: m=1 heap + direct loop 0.122 ms 6.44 GFLOP/s <- what we do now m=1 heap + native FFM kernel 1.002 ms 0.78 GFLOP/s <- copy-bound m=1 off-heap, zero copy 0.943 ms 0.83 GFLOP/s <- kernel-bound The native kernel copies both operands off-heap per call because a JDK 21 downcall cannot address heap memory, and that copy is proportional to the weight rather than to the work: 1.5 MB whatever m is. Keeping the weight off-heap instead — it is read-only for the life of the process, so this is the obvious escape — removes the copy and is still slower, because the segment kernel that then serves it is weaker than the heap one by more than the copy costs. So residency is not the lever. There are two independent gaps: this kernel cannot see heap memory (fixed on JDK 22+ by Linker.Option.critical(true), noted at the call site), and matmulFloatBlockedMemSeg is slow. Either is worth real throughput — the packed Q4_K path reaches ~29 GFLOP/s on the same machine. Also recorded: two things that measured as no-ops and should not be retried — reusing the Arena and its segments across calls (1.025 vs 1.002 ms at m=1; the allocation was never the cost) and reordering the C loops (already i-p-j). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../exec/tensor/ops/Fp32GemvShapeBench.kt | 25 ++++++++++++ .../exec/kernel/NativeFp32MatmulKernel.kt | 38 +++++++++++++++++++ 2 files changed, 63 insertions(+) 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 index 8bdedc4c0..b7a459f83 100644 --- 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 @@ -54,5 +54,30 @@ class Fp32GemvShapeBench { ) } 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-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 { From b840530e64bba1835c2420a0042f5ab5ccda463c Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 11:34:54 +0200 Subject: [PATCH 6/8] test(kernel): benchmark the FP32/BF16/FP16 matmul kernels at decode shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three benchmarks that turn 'the dense path is slow' into per-kernel numbers. m=1, k=1536, n=256 (a Gemma 4 E2B per-layer-embedding projection), Apple M4: native-ffm bf16 0.055 ms 14.33 GFLOP/s <- fastest kotlin-direct fp32 0.073 ms 10.78 GFLOP/s <- what DefaultCpuOpsJvm now uses native-ffm fp32 0.105 ms 7.48 GFLOP/s native-ffm fp16 0.121 ms 6.51 GFLOP/s scalar fp32 0.253 ms 3.10 GFLOP/s panama bf16 0.466 ms 1.69 GFLOP/s scalar fp16 0.886 ms 0.89 GFLOP/s panama fp32 0.969 ms 0.81 GFLOP/s Corrections to earlier guesses in this area, all measured: - The native FFM kernel is NOT the slow one. Its heap->off-heap copy, the obvious suspect, is 0.029 ms for the same 1.5 MB (54 GB/s, HeapToSegmentCopyBench) — 3% of a call, not the cost. - PanamaVectorMatmulKernel is slower than the SCALAR kernel at this shape, by 3.8x (0.969 vs 0.253 ms). It is the fallback wherever the native library is absent, so that is a live gap, not a curiosity. - Narrower weight storage does help, but BF16 not FP16: bf16 is 1.9x faster than fp32 (halved weight traffic, decode is a bit-shift), while fp16 is *slower* than fp32 because binary16 needs exponent rebiasing — exactly what the SPI doc says. Accumulation stays FP32 in both, so this is memory traffic, not math precision. - Off-heap residency does not help either: the segment kernel that serves it is slower than the heap one by more than the copy costs (0.943 vs 0.122 ms). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../ainet/exec/kernel/Fp32KernelRaceBench.kt | 118 ++++++++++++++++++ .../exec/kernel/HeapToSegmentCopyBench.kt | 72 +++++++++++ 2 files changed, 190 insertions(+) create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Fp32KernelRaceBench.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/HeapToSegmentCopyBench.kt 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..92f206c75 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/Fp32KernelRaceBench.kt @@ -0,0 +1,118 @@ +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 m = 1 + val k = 1536 + val n = 256 + val iterations = 300 + 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 providers: " + KernelRegistry.providers().joinToString { + "${it.name}(priority=${it.priority}, available=${it.isAvailable()})" + }) + + 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, + ) + ) + + // Narrow-float B: same problem, weight stored at 2 bytes/element instead of 4. Accumulation + // stays FP32 by contract, so this is a memory-traffic change, not a precision-of-math one — + // and at m=1 the weight read IS the work. BF16 decodes by a bit-shift (it is the top half of + // an FP32); FP16 needs exponent rebiasing, which the SPI docs call out as the slower decode. + val bBf16 = ByteArray(k * n * 2) + val bFp16 = ByteArray(k * n * 2) + for (i in 0 until k * n) { + val bits = java.lang.Float.floatToRawIntBits(b[i]) + val bf = (bits ushr 16).toShort() + bBf16[i * 2] = (bf.toInt() and 0xFF).toByte() + bBf16[i * 2 + 1] = ((bf.toInt() shr 8) and 0xFF).toByte() + val h = java.lang.Float.floatToFloat16(b[i]) + bFp16[i * 2] = (h.toInt() and 0xFF).toByte() + bFp16[i * 2 + 1] = ((h.toInt() shr 8) and 0xFF).toByte() + } + for (p: KernelProvider in KernelRegistry.providers()) { + if (!p.isAvailable()) continue + for ((label, kern, payload) in listOf( + Triple("bf16", p.matmulBf16(), bBf16), + Triple("fp16", p.matmulFp16(), bFp16), + )) { + if (kern == null) { println("RACE %-16s %s: ".format(p.name, label)); continue } + repeat(30) { kern.matmul(a, 0, k, payload, 0, n * 2, out, 0, n, m, n, k) } + val e = measureTime { + repeat(iterations) { kern.matmul(a, 0, k, payload, 0, n * 2, out, 0, n, m, n, k) } + } + println( + "RACE %-16s %s %7.3f ms/call %6.2f GFLOP/s [%s]".format( + p.name, label, e.inWholeMicroseconds / 1000.0 / iterations, + 2.0 * macs / e.inWholeNanoseconds, kern::class.simpleName, + ) + ) + } + } + } + + 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) + } + } +} From a63d0e21e7975e62ab3d0ed1ccbb63906e75c52c Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 11:41:38 +0200 Subject: [PATCH 7/8] =?UTF-8?q?perf(kernel):=20PanamaVectorMatmulKernel=20?= =?UTF-8?q?was=20slower=20than=20scalar=20at=20m=3D1=20=E2=80=94=20add=20a?= =?UTF-8?q?=20GEMV=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tiled path has two fixed costs, both O(n*k) and both independent of m: it packs B transposed into a fresh n*k FloatArray on every call (a scatter, 1.5 MB for the shape below), and it does a horizontal reduceLanes per output cell per K-tile. Neither buys anything when there are no rows to amortize them over — and a decode step is m=1 by construction. Measured at k=1536 n=256 on an Apple M4, this kernel was losing to the scalar one it exists to beat: 0.969 ms vs 0.253. At or below 8 rows it now accumulates by outer product instead — out[i,:] += a[i,p] * b[p,:] — which touches B once, contiguously, keeps each lane on its own output column for the whole k loop, and never reduces across lanes: m before 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, unchanged) 25.53 3.06 32 1.647 ms 15.28 (tiled, unchanged) 27.69 3.10 16.7x at m=1, and it becomes the fastest FP32 kernel there — ahead of native-ffm, whose heap->off-heap copy costs more than the arithmetic at that shape. The tiled path is untouched above 8 rows, where it is what the shape wants. This matters most where there is no native library: Panama is then the best available provider, so every platform without the FFM backend was running its dense FP32 matmuls ~17x slower than necessary. 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; the native/Panama parity suite and the downstream Gemma 4 and FunctionGemma gates all pass. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../exec/kernel/PanamaVectorMatmulKernel.kt | 76 +++++++++++++++++++ .../ainet/exec/kernel/Fp32KernelRaceBench.kt | 43 ++--------- 2 files changed, 81 insertions(+), 38 deletions(-) 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-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 index 92f206c75..301bb438f 100644 --- 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 @@ -28,18 +28,19 @@ class Fp32KernelRaceBench { } if (KernelRegistry.providers().isEmpty()) KernelServiceLoader.installAll() - val m = 1 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 providers: " + KernelRegistry.providers().joinToString { - "${it.name}(priority=${it.priority}, available=${it.isAvailable()})" - }) + println("RACE ---- m=$m k=$k n=$n ----") for (p: KernelProvider in KernelRegistry.providers()) { if (!p.isAvailable()) continue @@ -66,40 +67,6 @@ class Fp32KernelRaceBench { ) ) - // Narrow-float B: same problem, weight stored at 2 bytes/element instead of 4. Accumulation - // stays FP32 by contract, so this is a memory-traffic change, not a precision-of-math one — - // and at m=1 the weight read IS the work. BF16 decodes by a bit-shift (it is the top half of - // an FP32); FP16 needs exponent rebiasing, which the SPI docs call out as the slower decode. - val bBf16 = ByteArray(k * n * 2) - val bFp16 = ByteArray(k * n * 2) - for (i in 0 until k * n) { - val bits = java.lang.Float.floatToRawIntBits(b[i]) - val bf = (bits ushr 16).toShort() - bBf16[i * 2] = (bf.toInt() and 0xFF).toByte() - bBf16[i * 2 + 1] = ((bf.toInt() shr 8) and 0xFF).toByte() - val h = java.lang.Float.floatToFloat16(b[i]) - bFp16[i * 2] = (h.toInt() and 0xFF).toByte() - bFp16[i * 2 + 1] = ((h.toInt() shr 8) and 0xFF).toByte() - } - for (p: KernelProvider in KernelRegistry.providers()) { - if (!p.isAvailable()) continue - for ((label, kern, payload) in listOf( - Triple("bf16", p.matmulBf16(), bBf16), - Triple("fp16", p.matmulFp16(), bFp16), - )) { - if (kern == null) { println("RACE %-16s %s: ".format(p.name, label)); continue } - repeat(30) { kern.matmul(a, 0, k, payload, 0, n * 2, out, 0, n, m, n, k) } - val e = measureTime { - repeat(iterations) { kern.matmul(a, 0, k, payload, 0, n * 2, out, 0, n, m, n, k) } - } - println( - "RACE %-16s %s %7.3f ms/call %6.2f GFLOP/s [%s]".format( - p.name, label, e.inWholeMicroseconds / 1000.0 / iterations, - 2.0 * macs / e.inWholeNanoseconds, kern::class.simpleName, - ) - ) - } - } } private fun directLoop(a: FloatArray, b: FloatArray, out: FloatArray, m: Int, n: Int, k: Int) { From 8f3fe167cb32e0c07f09822f189fe98492e1f7a8 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 31 Aug 2026 12:20:52 +0200 Subject: [PATCH 8/8] fix(build): androidMain cannot see KernelServiceLoader; keep the transpose cache private MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two CI breaks from the ViewKernelPack commit: - :skainet-backend-api:compileAndroidMain failed because the Android actual delegated to KernelServiceLoader, which lives in jvmMain and is not visible from androidMain. Inlines the same two steps it performs — ServiceLoader.load, then KernelRegistry.register, which sorts by priority on insertion. - jvmApiCheck failed because transposedDenseWeight was declared protected, putting an implementation detail into skainet-backend-cpu's public API. It is only called from within DefaultCpuOps, so it is private now and the .api dump is unchanged. Verified locally: compileAndroidMain, apiCheck across the repo, and the backend-api / backend-cpu / backend-native-cpu jvm suites. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YKeDSK4JF295y53Uvez954 --- .../backend/api/kernel/ViewKernelPack.android.kt | 11 ++++++++++- .../kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 2 +- 2 files changed, 11 insertions(+), 2 deletions(-) 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 index 2f60a006c..df7edc218 100644 --- 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 @@ -16,6 +16,15 @@ internal actual fun installPlatformKernelPacks(): List = .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 { KernelServiceLoader.installAll() }.getOrElse { emptyList() } + 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-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 e3e55eeb6..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 @@ -961,7 +961,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory * decode loop reuses; anything else falls through to a fresh transpose. */ @Suppress("UNCHECKED_CAST") - protected fun transposedDenseWeight(weight: Tensor): Tensor? { + 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 —