Gemma 4 engine gaps: gemma4 registries, KernelDispatch self-heal, and dense FP32 matmul (fixes #1220)

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]

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]

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()<br/>bootstrap if the table is empty"]
+  B --> C["normalizeActivation(a)<br/>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<br/>another activation format?"}
+  G -->|"yes, e.g. ternary wants int8"| H["requantize into caller's Scope<br/>emit AdapterInserted"]
+  H --> I{"kernel for the<br/>requantized pair?"}
+  I -->|yes| F
+  G -->|no| J{"prepackWeights = true<br/>and weight is ROW_MAJOR?"}
+  I -->|no| J
+  J -->|yes| K["prepack to INPUT_BLOCK_MAJOR<br/>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<br/>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 <name>: <reason>`. 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.

skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/ViewKernelPack.android.kt

@@ -0,0 +1,30 @@
+package sk.ainet.backend.api.kernel
+
+import java.util.ServiceLoader
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+
+/**
+ * Android discovery for [ViewKernelPack]. `ServiceLoader` exists on Android, so the JNI packs a
+ * consumer ships (e.g. the NEON row-major pack in `skainet-backend-jni-cpu`) are discovered the
+ * same way as on the JVM, provided the packaging step keeps `META-INF/services` entries.
+ */
+@ExperimentalMemoryApi
+internal actual fun installPlatformKernelPacks(): List<String> =
+    runCatching {
+        ServiceLoader.load(ViewKernelPack::class.java)
+            .mapNotNull { pack -> runCatching { pack.install(); pack.name }.getOrNull() }
+            .toList()
+    }.getOrElse { emptyList() }
+
+/**
+ * Provider discovery, inlined rather than delegated to `KernelServiceLoader`: that object lives in
+ * `jvmMain`, which the Android source set does not see. Same two steps it performs — discover, then
+ * register, letting [KernelRegistry] sort by priority on insertion.
+ */
+@ExperimentalMemoryApi
+internal actual fun installPlatformKernelProviders(): List<String> =
+    runCatching {
+        ServiceLoader.load(KernelProvider::class.java)
+            .mapNotNull { provider -> runCatching { KernelRegistry.register(provider); provider.name }.getOrNull() }
+            .toList()
+    }.getOrElse { emptyList() }

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<ViewKernel> = mutableListOf()
 
+    private var autoInstallAttempted: Boolean = false
+
+    /**
+     * Populate the table from platform-discovered providers and [ViewKernelPack]s, once per
+     * process, when nothing has been registered yet.
+     *
+     * [KernelRegistry] has always self-healed this way (`DefaultCpuOpsJvm.ensureKernelProviders`
+     * installs providers on first use); this dispatcher did not, so every consumer had to remember
+     * an explicit bootstrap before its first forward pass. Forgetting it is silent — dispatch
+     * simply falls to the decoding reference kernel, which is correct and about a thousand times
+     * slower — and it was forgotten repeatedly in practice, by application entry points and
+     * diagnostic harnesses alike.
+     *
+     * Order matters: providers first, because [KernelPacks.install] derives its kernels from
+     * `KernelRegistry.bestAvailable()` and would otherwise contribute nothing but the reference
+     * kernel.
+     *
+     * Explicit installation still works and still wins — a consumer that registers its own kernels
+     * before the first dispatch suppresses auto-install entirely, and later registrations override
+     * earlier ones for the same key. Call [clearForTesting] to re-arm.
+     */
+    public fun ensureInstalled() {
+        if (autoInstallAttempted || kernels.isNotEmpty()) return
+        autoInstallAttempted = true
+        if (KernelRegistry.providers().isEmpty()) installPlatformKernelProviders()
+        KernelPacks.install()
+        installPlatformKernelPacks()
+    }
+
     /** Register [kernel]; later registrations win for the same key (a pack can override the reference). */
     public fun register(kernel: ViewKernel) {
         kernels.removeAll { it.key == kernel.key && it.name == kernel.name }
@@ -40,7 +69,10 @@ public object KernelDispatch {
     /** The kernel registered for [key], or `null`. */
     public fun find(key: KernelKey): ViewKernel? = kernels.firstOrNull { it.key == key }
 
-    public fun clearForTesting() { kernels.clear() }
+    public fun clearForTesting() {
+        kernels.clear()
+        autoInstallAttempted = false
+    }
 
     /**
      * Encodings a [MappedCapableKernel] registered right now serves as a `BLOCKED_ROW_MAJOR`
@@ -76,6 +108,16 @@ public object KernelDispatch {
         }
     }
 
+    /**
+     * Process-global default [TraceSink] used when a call site does not pass one. Production
+     * call sites (e.g. `DefaultCpuOps`) rely on the parameter default, which made every
+     * reference-kernel fallback invisible — set this (e.g. from a diagnostic harness) to
+     * observe dispatch decisions everywhere without threading a sink through the ops layer.
+     */
+    public var defaultSink: TraceSink = NoopTraceSink
+
+    private var warnedReferenceFallback: Boolean = false
+
     /**
      * Select and run `matmul(a, b)`, writing into [out]. [scope] owns any adapter the selection
      * needs; [sink] sees the kernel run and every adapter.
@@ -87,7 +129,7 @@ public object KernelDispatch {
         b: TensorView,
         out: TensorView,
         scope: Scope = Scope.Ambient,
-        sink: TraceSink = NoopTraceSink,
+        sink: TraceSink = defaultSink,
         /**
          * Relayout a canonical packed weight into kernel order when that is what unlocks a packed
          * kernel (#973/#1095).
@@ -101,6 +143,9 @@ public object KernelDispatch {
          */
         prepackWeights: Boolean = false,
     ) {
+        // Self-heal on first use: an empty table means nobody bootstrapped, and the silent
+        // consequence is the reference kernel for every operand pair.
+        ensureInstalled()
         val key = KernelKey.matmul(a, b)
         val exact = find(key)
         if (exact != null) {
@@ -137,6 +182,21 @@ public object KernelDispatch {
         }
         // No exact kernel: adapt the operands a kernel would accept, then fall back to the reference,
         // which reads any format through decoding get().
+        // The reference path is correct but orders of magnitude slower than a real kernel on a
+        // blocked weight (per-element block decode) — a process that lands here on a quantized
+        // weight almost certainly forgot to install a kernel pack. Say so once, loudly, even with
+        // no sink attached: silent fallback is how a 25 s/token regression ships unnoticed.
+        if (!warnedReferenceFallback && b.layout.blocked) {
+            warnedReferenceFallback = true
+            println(
+                "[SKaiNET] KernelDispatch: no kernel registered for matmul " +
+                    "(activation=${a.format.encoding}, weight=${b.format.encoding}, " +
+                    "order=${b.layout.blockOrder}); falling back to the decoding reference " +
+                    "kernel (~1000x slower). Install a kernel pack (e.g. KernelPacks.install() " +
+                    "+ FfmRowMajorKernelPack.install()) before the first forward. " +
+                    "Further fallbacks are not reported."
+            )
+        }
         val adaptedA = adapt(a, scope, sink, "gather")
         val reference = ReferenceMatmulKernel(KernelKey.matmul(adaptedA, b))
         runTraced(reference, listOf(adaptedA, b), out, sink)