From c570a3498d5c6e8bd9606a19c75dbc0f5ffe9ff8 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 20:08:44 +0200 Subject: [PATCH] feat(backend-cpu): route the generic matmul path through the kernel registry (SKEEP-003 P3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F5 / M1-A4. The fast paths are untouched: the packed-quant kernels and the FP32 2D×2D path run exactly as before. What changes is the *fallback* — the path where #993 crashed: - DefaultCpuOps.matmul now tries dispatchMatmulViaRegistry() before matmulGeneric(). It builds views from both operands (TensorData.view, #1068/#1069), transposes the weight as a view, normalises the activation once (rank-1 -> [1, k], batched -> [rows, k]) and hands the key to KernelDispatch, which selects a registered kernel or the decoding reference kernel. A packed weight is therefore decoded, never read as a raw byte, whatever the activation's rank or subtype. - The result's shape is restored from the normalisation ([k] x [k, n] -> [n], batched dims preserved), so callers see no change. - DispatchMode (backend-api, expect/actual): skainet.dispatch.registry =false forces the legacy per-element fallback, and overrideEnabled lets a test pin either path. The legacy code stays until the migration is complete, then it goes. - RegistryMatmulDispatchTest: the #993 repro (rank-1 decode step against a Q8_0 weight) is correct and finite; registry and legacy paths agree elementwise on the same inputs; batched activations flatten and reshape; a Q4_K weight decodes correctly. 4/4. Closes #1028 Co-Authored-By: Claude Fable 5 --- .../api/kernel/DispatchMode.android.kt | 6 + .../ainet/backend/api/kernel/DispatchMode.kt | 28 +++++ .../backend/api/kernel/DispatchMode.other.kt | 7 ++ .../backend/api/kernel/DispatchMode.jvm.kt | 6 + .../backend/api/kernel/DispatchMode.other.kt | 7 ++ .../backend/api/kernel/DispatchMode.other.kt | 7 ++ .../backend/api/kernel/DispatchMode.other.kt | 7 ++ .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 49 +++++++- .../tensor/ops/RegistryMatmulDispatchTest.kt | 115 ++++++++++++++++++ 9 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.android.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.kt create mode 100644 skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt create mode 100644 skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.jvm.kt create mode 100644 skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt create mode 100644 skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt create mode 100644 skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt diff --git a/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.android.kt b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.android.kt new file mode 100644 index 000000000..5e510782d --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/androidMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.android.kt @@ -0,0 +1,6 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +@ExperimentalMemoryApi +internal actual fun platformUseRegistry(): Boolean = System.getProperty(DispatchMode.PROPERTY) != "false" diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.kt new file mode 100644 index 000000000..1e8892fcc --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.kt @@ -0,0 +1,28 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * Whether the generic (non-fast-path) matmul goes through [KernelDispatch] or the legacy + * per-element fallback (SKEEP-003 §5.1, migration slice #1028). + * + * The registry path is the default: it normalises rank once as views and reads packed operands + * through decoding `get()`, so a rank-1 decode step against a packed weight is *correct by + * construction* instead of a `ClassCastException` (#993). The legacy path stays one flag away + * while the migration settles; it is deleted once the golden parity and benchmark evidence is in. + */ +@ExperimentalMemoryApi +public object DispatchMode { + /** Set to `false` (`skainet.dispatch.registry=false`) to force the legacy generic fallback. */ + public const val PROPERTY: String = "skainet.dispatch.registry" + + /** Overridable in tests; `null` means "read the platform setting". */ + public var overrideEnabled: Boolean? = null + + /** Whether the generic path should use the kernel registry. */ + public fun useRegistry(): Boolean = overrideEnabled ?: platformUseRegistry() +} + +/** Platform reading of [DispatchMode.PROPERTY]; defaults to `true` where there is no property store. */ +@ExperimentalMemoryApi +internal expect fun platformUseRegistry(): Boolean diff --git a/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt b/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt new file mode 100644 index 000000000..94d26a87c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt @@ -0,0 +1,7 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** No system properties here: the registry path is always on (override in tests via [DispatchMode.overrideEnabled]). */ +@ExperimentalMemoryApi +internal actual fun platformUseRegistry(): Boolean = true diff --git a/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.jvm.kt b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.jvm.kt new file mode 100644 index 000000000..5e510782d --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/jvmMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.jvm.kt @@ -0,0 +1,6 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +@ExperimentalMemoryApi +internal actual fun platformUseRegistry(): Boolean = System.getProperty(DispatchMode.PROPERTY) != "false" diff --git a/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt b/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt new file mode 100644 index 000000000..94d26a87c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/nativeMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt @@ -0,0 +1,7 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** No system properties here: the registry path is always on (override in tests via [DispatchMode.overrideEnabled]). */ +@ExperimentalMemoryApi +internal actual fun platformUseRegistry(): Boolean = true diff --git a/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt b/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt new file mode 100644 index 000000000..94d26a87c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/wasmJsMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt @@ -0,0 +1,7 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** No system properties here: the registry path is always on (override in tests via [DispatchMode.overrideEnabled]). */ +@ExperimentalMemoryApi +internal actual fun platformUseRegistry(): Boolean = true diff --git a/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt b/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt new file mode 100644 index 000000000..94d26a87c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/wasmWasiMain/kotlin/sk/ainet/backend/api/kernel/DispatchMode.other.kt @@ -0,0 +1,7 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** No system properties here: the registry path is always on (override in tests via [DispatchMode.overrideEnabled]). */ +@ExperimentalMemoryApi +internal actual fun platformUseRegistry(): Boolean = true 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 27107a5e6..cc5e11983 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 @@ -627,8 +627,53 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory } } - // Generic fallback for batched / non-float / non-2D cases - return KernelProfile.timeGeneric { matmulGeneric(a, b) } + // Everything else: the kernel registry first (SKEEP-003 §5.1) — rank is normalised once as + // views and packed operands are read through decoding get(), so a rank-1 decode step against + // a packed weight is correct by construction rather than a ClassCastException (#993). The + // legacy per-element fallback stays one flag away (skainet.dispatch.registry=false) until the + // migration is complete. + return KernelProfile.timeGeneric { + dispatchMatmulViaRegistry(a, b) ?: matmulGeneric(a, b) + } + } + + /** + * Run `matmul` through [KernelDispatch] when both operands can describe themselves as views and + * the result is float-typed; `null` means "not expressible here, use the legacy path". + */ + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + @Suppress("UNCHECKED_CAST") + private fun dispatchMatmulViaRegistry(a: Tensor, b: Tensor): Tensor? { + if (!sk.ainet.backend.api.kernel.DispatchMode.useRegistry()) return null + if (a.dtype != FP32::class) return null + if (b.shape.rank != 2) return null + val aView = a.data.view ?: return null + val bView = b.data.view ?: return null + // b is [k, n] here; the kernels take the weight output-major, which is a transposed *view*. + val bT = try { bView.transpose() } catch (_: IllegalArgumentException) { return null } + val (aNorm, leading) = try { + sk.ainet.backend.api.kernel.KernelDispatch.normalizeActivation(aView) + } catch (_: IllegalArgumentException) { + return null + } + val m = aNorm.shape[0] + val k = aNorm.shape[1] + val n = bT.shape[0] + if (k != bT.shape[1]) return null + val outArray = FloatArray(m * n) + val outView = sk.ainet.lang.memory.TensorView.dense( + sk.ainet.lang.memory.Storage.Heap.wrap(outArray), + Shape(m, n), + FP32, + ) + sk.ainet.backend.api.kernel.KernelDispatch.matmul(aNorm, bT, outView) + val outShape = when { + a.shape.rank == 1 -> Shape(n) // [k] x [k, n] -> [n] + leading.isEmpty() -> Shape(m, n) + else -> Shape(*(leading + n)) + } + val outData = dataFactory.fromFloatArray(outShape, a.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData + return newTensor(outData, a.dtype, a, b) } private fun matmulGeneric( diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt new file mode 100644 index 000000000..c553dc029 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt @@ -0,0 +1,115 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.backend.api.kernel.DispatchMode +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Q4_KBlockTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * SKEEP-003 §5.1 / PRD M1-F5, M1-A4: the generic matmul path now goes through the kernel registry. + * The two crash classes this replaces are #993 (a rank-1 decode-step activation reaching a kernel + * written for rank 2, then `matmulGeneric` reading a packed weight's raw byte) and #991 (an + * activation whose `TensorData` subtype the fast path did not accept). Both must produce correct, + * finite numbers — and the registry path and the legacy path must agree. + */ +@OptIn(ExperimentalMemoryApi::class) +class RegistryMatmulDispatchTest { + + @AfterTest fun reset() { DispatchMode.overrideEnabled = null } + + private val ctx = DirectCpuExecutionContext() + + private fun half(v: Float): Int { + val b = v.toRawBits(); val sign = (b ushr 16) and 0x8000 + val e = ((b ushr 23) and 0xFF) - 127 + 15; val m = b and 0x7FFFFF + if (e <= 0) return sign; if (e >= 31) return sign or 0x7C00 + return sign or (e shl 10) or (m ushr 13) + } + + /** A Q8_0 weight stored as [k, n] = [32, rows] would be column-major; SKaiNET stores [out, in] and transposes. */ + private fun q8Bytes(rows: Int): ByteArray { + val bytes = ByteArray(rows * 34) + for (r in 0 until rows) { + val off = r * 34; val d = half(0.25f) + bytes[off] = (d and 0xFF).toByte(); bytes[off + 1] = ((d ushr 8) and 0xFF).toByte() + for (i in 0 until 32) bytes[off + 2 + i] = ((i - 16) + r * 3).toByte() + } + return bytes + } + + @Suppress("UNCHECKED_CAST") + private fun packedWeight(rows: Int): Pair, FloatArray> { + val bytes = q8Bytes(rows) + val data = Q8_0BlockTensorData(Shape(rows, 32), bytes) + val t = ctx.fromData(data as TensorData, FP32::class) + return t to data.toFloatArray() + } + + @Test + fun rank1DecodeStepAgainstAPackedWeightIsCorrect() { + // #993: the first post-prefill decode step passes a rank-1 [hidden] activation + val (w, wf) = packedWeight(4) + val x = ctx.fromFloatArray(Shape(32), FP32::class, FloatArray(32) { (it % 7) * 0.5f }) + val out = ctx.ops.matmul(x, ctx.ops.transpose(w)) // [32] x [32, 4] -> [4] + val got = out.data.copyToFloatArray() + assertTrue(out.shape.rank == 1 && out.shape[0] == 4, "expected [4], was ${out.shape}") + for (j in 0 until 4) { + var expect = 0f + for (t in 0 until 32) expect += ((t % 7) * 0.5f) * wf[j * 32 + t] + assertTrue(got[j].isFinite(), "row $j is not finite") + assertTrue(abs(got[j] - expect) < 1e-2f, "row $j: ${got[j]} vs $expect") + } + } + + @Test + fun registryAndLegacyPathsAgreeOnAPackedWeight() { + val (w, _) = packedWeight(3) + val x = ctx.fromFloatArray(Shape(2, 32), FP32::class, FloatArray(64) { (it % 5) * 0.25f }) + val wt = ctx.ops.transpose(w) + + DispatchMode.overrideEnabled = true + val viaRegistry = ctx.ops.matmul(x, wt).data.copyToFloatArray() + DispatchMode.overrideEnabled = false + val viaLegacy = ctx.ops.matmul(x, wt).data.copyToFloatArray() + + for (i in viaRegistry.indices) { + assertTrue(viaRegistry[i].isFinite() && viaLegacy[i].isFinite()) + assertTrue(abs(viaRegistry[i] - viaLegacy[i]) < 1e-3f, "element $i: registry ${viaRegistry[i]} vs legacy ${viaLegacy[i]}") + } + } + + @Test + fun batchedActivationsFlattenAndReshape() { + val (w, wf) = packedWeight(2) + val x = ctx.fromFloatArray(Shape(2, 3, 32), FP32::class, FloatArray(192) { (it % 4).toFloat() }) + val out = ctx.ops.matmul(x, ctx.ops.transpose(w)) + assertTrue(out.shape.dimensions.toList() == listOf(2, 3, 2), "expected [2, 3, 2], was ${out.shape}") + val got = out.data.copyToFloatArray() + var expect = 0f + for (t in 0 until 32) expect += ((t % 4).toFloat()) * wf[t] + assertTrue(abs(got[0] - expect) < 1e-2f, "${got[0]} vs $expect") + } + + @Test + fun q4kWeightsDecodeThroughTheRegistry() { + val bytes = ByteArray(144) { (it * 11).toByte() } + val d = half(0.01f); bytes[0] = (d and 0xFF).toByte(); bytes[1] = ((d ushr 8) and 0xFF).toByte() + val dm = half(0.005f); bytes[2] = (dm and 0xFF).toByte(); bytes[3] = ((dm ushr 8) and 0xFF).toByte() + val data = Q4_KBlockTensorData(Shape(1, 256), bytes) + @Suppress("UNCHECKED_CAST") + val w = ctx.fromData(data as TensorData, FP32::class) + val x = ctx.fromFloatArray(Shape(256), FP32::class, FloatArray(256) { 0.125f }) + val got = ctx.ops.matmul(x, ctx.ops.transpose(w)).data.copyToFloatArray() + var expect = 0f + for (t in 0 until 256) expect += 0.125f * data.toFloatArray()[t] + assertTrue(got[0].isFinite()); assertTrue(abs(got[0] - expect) < 1e-2f, "${got[0]} vs $expect") + } +}