From 07796c5b74e97c60cd23fd27f749bf121069c67b Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 19:08:45 +0200 Subject: [PATCH] =?UTF-8?q?feat(backend-api):=20KernelKey=20dispatch=20?= =?UTF-8?q?=E2=80=94=20declared=20formats=20and=20layouts,=20rank=20normal?= =?UTF-8?q?ised=20once,=20visible=20adapters,=20reference=20matmul=20(SKEE?= =?UTF-8?q?P-003=20P3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F5 / M1-A4. SKEEP-003 §5.1: kernel selection keys on what an operand *declares* instead of an is-ladder over TensorData subclasses — the reason every quantisation bug was a dispatch bug (#993, #991). - KernelKey(op, operands, placement) with OperandKey(format, layoutClass) and LayoutClass { CONTIGUOUS, STRIDED, BLOCKED }; keys are values and print as "matmul(Float32/Dense(4B) contiguous × Float32/Q8_0 blocked) @host". UnsupportedKernelException lists what is registered. - ViewKernel: a kernel takes TensorViews and declares its key; a custom kernel author never sees a TensorData subclass. - ReferenceMatmulKernel: correct for *any* pair of formats because it reads through TensorView.get(), which decodes (rule 4). Slow by design — the fallback that turns an unsupported combination into right numbers instead of a ClassCastException in layer 17. - KernelDispatch: normalizeActivation() promotes rank-1 to [1, k] and flattens leading dims as *views* (once, before lookup — #993's root cause disappears), find/register keep the kernel table, matmul() selects, inserts a gather adapter for a strided operand into the caller's Scope and emits TraceEvent.AdapterInserted (the hidden 12 GB of #782 becomes visible), then runs the kernel inside a TraceEvent.KernelRun span. TensorView.reshapeContiguous is a view. - KernelKeyDispatchTest: keys describe formats/layouts, rank normalisation is zero-copy, the #993 repro (rank-1 activation × Q8_0 weight) and a Q4_K case produce finite, numerically correct output through the registry, a registered kernel wins over the reference, and a strided activation gets a visible gather adapter whose numbers match a manual dot product. JVM 6/6, linuxX64 6/6. - backend-api wires kotlin-test into commonTest (it declares its targets by hand, without the sk.ainet.multiplatform convention plugin). Nothing routes DefaultCpuOps through this yet — that is #1028 (common) and #1029 (JVM packs), where the ladders are replaced under the golden parity gate. Closes #1027 Co-Authored-By: Claude Fable 5 --- .../skainet-backend-api/build.gradle.kts | 6 + .../backend/api/kernel/KernelDispatch.kt | 115 ++++++++++++ .../sk/ainet/backend/api/kernel/KernelKey.kt | 80 ++++++++ .../ainet/backend/api/kernel/MatmulKernels.kt | 59 ++++++ .../api/kernel/KernelKeyDispatchTest.kt | 177 ++++++++++++++++++ 5 files changed, 437 insertions(+) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelKey.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MatmulKernels.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/KernelKeyDispatchTest.kt diff --git a/skainet-backends/skainet-backend-api/build.gradle.kts b/skainet-backends/skainet-backend-api/build.gradle.kts index 601c87edb..e37a2ea34 100644 --- a/skainet-backends/skainet-backend-api/build.gradle.kts +++ b/skainet-backends/skainet-backend-api/build.gradle.kts @@ -42,6 +42,12 @@ kotlin { } sourceSets { + commonTest.dependencies { + // the module declares its targets by hand (no sk.ainet.multiplatform convention plugin), + // so kotlin-test is wired here + implementation(libs.kotlin.test) + } + commonMain.dependencies { // Neutral backend API is an `api` re-export of the tensor op and // storage interfaces already defined in skainet-lang-core. Any 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 new file mode 100644 index 000000000..03e9cc038 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt @@ -0,0 +1,115 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import sk.ainet.lang.memory.trace.kernel as traceKernel +import sk.ainet.lang.tensor.Shape + +/** + * Kernel selection on declared descriptors instead of an `is`-ladder over Kotlin classes + * (SKEEP-003 §5.1). The order is: **normalize** the operands as zero-copy views (so rank-1 decode + * steps never reach a kernel written for rank 2 — the #993 root cause disappears), build the + * [KernelKey], look it up, insert **visible** adapters when a kernel cannot take an operand as it + * is, and fall back to the reference kernel, which is correct for every format because it decodes. + * + * Adapters allocate in the caller's [Scope] (a `Forward` scope in a generation loop) and are + * emitted as [TraceEvent.AdapterInserted] — the "hidden 12 GB" of #782 becomes a visible event. + */ +@ExperimentalMemoryApi +public object KernelDispatch { + + private val kernels: MutableList = mutableListOf() + + /** 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 } + kernels.add(0, kernel) + } + + /** Every registered kernel, most recently registered first. */ + public fun kernels(): List = kernels.toList() + + /** The kernel registered for [key], or `null`. */ + public fun find(key: KernelKey): ViewKernel? = kernels.firstOrNull { it.key == key } + + public fun clearForTesting() { kernels.clear() } + + /** + * Normalize a matmul operand pair to rank 2 as **views** (rule 5, §5.1 "rank handling happens + * once"): `[k]` becomes `[1, k]`, `[b, s, k]` becomes `[b*s, k]` when contiguous. Returns the + * normalized activation and the number of leading dims that were flattened, so the caller can + * reshape the result back. + */ + public fun normalizeActivation(a: TensorView): Pair = when { + a.shape.rank == 1 -> a.unsqueeze(0) to intArrayOf() + a.shape.rank == 2 -> a to intArrayOf() + else -> { + val leading = IntArray(a.shape.rank - 1) { a.shape[it] } + require(a.isContiguous) { "flattening leading dims needs a contiguous activation; materialize first" } + var rows = 1 + for (d in leading) rows *= d + a.reshapeContiguous(Shape(rows, a.shape[a.shape.rank - 1])) to leading + } + } + + /** + * Select and run `matmul(a, b)`, writing into [out]. [scope] owns any adapter the selection + * needs; [sink] sees the kernel run and every adapter. + * + * @throws UnsupportedKernelException when neither a kernel nor the reference path can serve the key + */ + public fun matmul( + a: TensorView, + b: TensorView, + out: TensorView, + scope: Scope = Scope.Ambient, + sink: TraceSink = NoopTraceSink, + ) { + val key = KernelKey.matmul(a, b) + val exact = find(key) + if (exact != null) { + runTraced(exact, listOf(a, b), out, sink) + return + } + // No exact kernel: adapt the operands a kernel would accept, then fall back to the reference, + // which reads any format through decoding get(). + val adaptedA = adapt(a, scope, sink, "gather") + val reference = ReferenceMatmulKernel(KernelKey.matmul(adaptedA, b)) + runTraced(reference, listOf(adaptedA, b), out, sink) + } + + /** Materialize [view] into a dense contiguous view when it is strided; emits an adapter event. */ + public fun adapt(view: TensorView, scope: Scope, sink: TraceSink, kind: String): TensorView { + if (view.isContiguous || view.layout.blocked) return view + val dense = view.materialize(Format.dense(view.format.dtype), scope) + if (sink.isEnabled) { + sink.emit(TraceEvent.AdapterInserted(kind, view.format, dense.format, dense.elementCount * view.format.dtype.sizeInBytes, view.id)) + } + return dense + } + + private fun runTraced(kernel: ViewKernel, inputs: List, out: TensorView, sink: TraceSink) { + if (!sink.isEnabled) { kernel.run(inputs, out); return } + sink.traceKernel( + op = kernel.key.op, + kernel = kernel.name, + inputs = inputs.map { it.id }, + output = out.id, + bytesRead = inputs.sumOf { it.elementCount * it.format.dtype.sizeInBytes }, + bytesWritten = out.elementCount * out.format.dtype.sizeInBytes, + ) { kernel.run(inputs, out) } + } +} + +/** A view of the same contiguous bytes under a different shape (rule 5: reshape is a view). */ +@ExperimentalMemoryApi +public fun TensorView.reshapeContiguous(newShape: Shape): TensorView { + require(isContiguous) { "reshape needs a contiguous view" } + require(newShape.volume.toLong() == elementCount) { "reshape must keep the element count ($elementCount), got ${newShape.volume}" } + return TensorView(newShape, format, sk.ainet.lang.memory.Layout.rowMajor(newShape, format, layout.offsetElements), storage, id) +} diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelKey.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelKey.kt new file mode 100644 index 000000000..3442aee2a --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelKey.kt @@ -0,0 +1,80 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * What a kernel consumes, declared rather than discovered (SKEEP-003 §0 *KernelKey*, §5.1): the op, + * the [Format] of each operand, how each operand is laid out, and the placement it needs. The + * dispatcher looks a key up instead of walking an `is`-ladder over `TensorData` subclasses — which + * is what made every quantisation bug a dispatch bug (#993, #991). + * + * Keys are values: equal keys select the same kernel, and a key prints as something a log or an + * `UnsupportedKernel` message can show: `matmul(F32/Dense(4B) contiguous × F32/Q4_K blocked) @host`. + */ +@ExperimentalMemoryApi +public data class KernelKey( + val op: String, + val operands: List, + val placement: Placement = Placement.HOST, +) { + /** Where the operands live — host memory today; a device backend adds its own (PRD non-goal for M1). */ + public enum class Placement { HOST, DEVICE } + + override fun toString(): String = + "$op(${operands.joinToString(" × ")})" + if (placement != Placement.HOST) " @${placement.name.lowercase()}" else " @host" + + public companion object { + /** The key of `matmul(activation, weight)` as the two views describe themselves. */ + public fun matmul(activation: TensorView, weight: TensorView, placement: Placement = Placement.HOST): KernelKey = + KernelKey("matmul", listOf(OperandKey.of(activation), OperandKey.of(weight)), placement) + } +} + +/** One operand of a [KernelKey]: its [Format] plus the layout class the kernel must cope with. */ +@ExperimentalMemoryApi +public data class OperandKey(val format: Format, val layout: LayoutClass) { + override fun toString(): String = "$format ${layout.name.lowercase()}" + + public companion object { + /** Describe [view]: dense-and-gap-free is `CONTIGUOUS`, a packed layout is `BLOCKED`, anything else `STRIDED`. */ + public fun of(view: TensorView): OperandKey { + val cls = when { + view.layout.blocked -> LayoutClass.BLOCKED + view.isContiguous -> LayoutClass.CONTIGUOUS + else -> LayoutClass.STRIDED + } + return OperandKey(view.format, cls) + } + + /** A dense contiguous operand of [format] — the shape kernels prefer. */ + public fun contiguous(format: Format): OperandKey = OperandKey(format, LayoutClass.CONTIGUOUS) + } +} + +/** + * How an operand's bytes are arranged, as far as kernel selection cares: one gap-free run + * ([CONTIGUOUS]), a strided view over a larger buffer ([STRIDED]), or block-packed ([BLOCKED]). + * A kernel that declares `CONTIGUOUS` gets a gather adapter inserted for a `STRIDED` operand + * (§5.1) — the adapter is visible in the trace, never hidden inside a kernel. + */ +@ExperimentalMemoryApi +public enum class LayoutClass { CONTIGUOUS, STRIDED, BLOCKED } + +/** Thrown when no registered kernel and no adapter chain can serve a key; lists what is registered. */ +@ExperimentalMemoryApi +public class UnsupportedKernelException( + public val key: KernelKey, + public val candidates: List, + message: String = "No kernel for $key" + if (candidates.isEmpty()) "" else "; registered: ${candidates.joinToString(", ")}", +) : IllegalArgumentException(message) + +/** The encoding name a [KernelKey] uses for a format, matching `KernelProvider.supports`' dtype keys. */ +@ExperimentalMemoryApi +public val Format.kernelEncodingName: String + get() = when (val e = encoding) { + is TensorEncoding.Dense -> dtype.name + else -> e.name + } diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MatmulKernels.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MatmulKernels.kt new file mode 100644 index 000000000..5aae52106 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/MatmulKernels.kt @@ -0,0 +1,59 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.TensorView + +/** + * A registered kernel: a function `(inputs, out) -> Unit` behind a [KernelKey] (SKEEP-003 §0 + * *Kernel*, §5.2). Custom kernels are written against views and registered — an author never + * touches a `TensorData` subclass. + */ +@ExperimentalMemoryApi +public interface ViewKernel { + /** What this kernel serves. */ + public val key: KernelKey + + /** A name for logs, traces and `UnsupportedKernel` messages (`scalar-reference`, `panama-q4k`, …). */ + public val name: String + + /** Run the kernel: [inputs] as described by [key], result written into [out]. */ + public fun run(inputs: List, out: TensorView) +} + +/** + * The reference matmul: correct for **any** pair of formats and layouts, because it reads through + * `TensorView.get()`, which decodes (rule 4). Slow by design — it is the fallback that makes an + * unsupported combination produce right numbers with a warning instead of a `ClassCastException` + * in layer 17 (#993). Registered for every key the dispatcher cannot serve better. + * + * `out = a × bᵀ` in the shapes SKaiNET's dispatch normalises to: `a` is `[m, k]`, `b` is `[n, k]` + * (a weight stored output-major, as GGUF does), `out` is `[m, n]`. + */ +@ExperimentalMemoryApi +public class ReferenceMatmulKernel(override val key: KernelKey) : ViewKernel { + override val name: String get() = "reference" + + override fun run(inputs: List, out: TensorView) { + require(inputs.size == 2) { "matmul takes two operands, got ${inputs.size}" } + val a = inputs[0]; val b = inputs[1] + require(a.shape.rank == 2 && b.shape.rank == 2 && out.shape.rank == 2) { "reference matmul works on rank-2 views (normalise first)" } + val m = a.shape[0]; val k = a.shape[1]; val n = b.shape[0] + require(b.shape[1] == k) { "inner dimensions disagree: a is [${m}, ${k}], b is [${n}, ${b.shape[1]}]" } + require(out.shape[0] == m && out.shape[1] == n) { "out must be [$m, $n], was ${out.shape}" } + for (i in 0 until m) { + for (j in 0 until n) { + var acc = 0f + for (t in 0 until k) acc += a.get(i, t) * b.get(j, t) + out.set(i, j, value = acc) + } + } + } + + public companion object { + /** The reference kernel for the formats of [a] and [b]. */ + @ExperimentalMemoryApi + public fun forOperands(a: Format, b: Format): ReferenceMatmulKernel = + ReferenceMatmulKernel(KernelKey("matmul", listOf(OperandKey.contiguous(a), OperandKey.contiguous(b)))) + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/KernelKeyDispatchTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/KernelKeyDispatchTest.kt new file mode 100644 index 000000000..b9799c89f --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/KernelKeyDispatchTest.kt @@ -0,0 +1,177 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ForwardScope +import sk.ainet.lang.memory.ScopeKind +import sk.ainet.lang.memory.PackedBlockDecoder +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.data.Q4_KBlockTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** + * SKEEP-003 §5.1 / PRD M1-F5, M1-A4: dispatch keys on declared formats and layouts, rank is + * normalised once before lookup, adapters are visible, and the reference kernel is correct for any + * format — so the #993 (rank-1 decode step × packed weight) and #991 (activation subtype) cases + * cannot crash. + */ +@OptIn(ExperimentalMemoryApi::class) +class KernelKeyDispatchTest { + + @AfterTest fun cleanup() { KernelDispatch.clearForTesting() } + + private fun denseView(shape: Shape, init: (Int) -> Float): TensorView { + val s = Storage.Heap.floats(shape.volume) + val f = s.floats!! + for (i in f.indices) f[i] = init(i) + return TensorView.dense(s, shape, FP32) + } + + 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 of [rows] × 32 with known values: code i - 16, scale 0.5. */ + private fun q8Weight(rows: Int): Pair { + val bytes = ByteArray(rows * 34) + for (r in 0 until rows) { + val off = r * 34; val d = half(0.5f) + 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).toByte() + } + val data = Q8_0BlockTensorData(Shape(rows, 32), bytes) + val view = TensorView.packed(Storage.Heap.wrap(bytes, mutable = false), Shape(rows, 32), TensorEncoding.Q8_0, PackedBlockDecoder(data), id = TensorId.parse("model.layers[0].attn.q_proj.weight")) + return view to data.toFloatArray() + } + + @Test + fun keysDescribeFormatsAndLayouts() { + val a = denseView(Shape(1, 32)) { it.toFloat() } + val (w, _) = q8Weight(2) + val key = KernelKey.matmul(a, w) + assertEquals("matmul", key.op); assertEquals(2, key.operands.size) + assertEquals(OperandKey(Format.dense(FP32), LayoutClass.CONTIGUOUS), key.operands[0]) + assertEquals(OperandKey(Format(FP32, TensorEncoding.Q8_0), LayoutClass.BLOCKED), key.operands[1]) + assertEquals("matmul(Float32/Dense(4B) contiguous × Float32/Q8_0 blocked) @host", key.toString()) + // a strided operand is a different key — that is the point of keying on layout + val strided = denseView(Shape(4, 8)) { it.toFloat() }.narrow(1, 0, 4) + assertEquals(LayoutClass.STRIDED, OperandKey.of(strided).layout) + assertEquals(Format(FP32, TensorEncoding.Q8_0).kernelEncodingName, "Q8_0") + assertEquals(Format.dense(FP32).kernelEncodingName, "Float32") + } + + @Test + fun rankIsNormalisedOnceBeforeLookup() { + // #993: a rank-1 decode-step activation must never reach a kernel written for rank 2 + val rank1 = denseView(Shape(32)) { 1f } + val (norm, leading) = KernelDispatch.normalizeActivation(rank1) + assertEquals(Shape(1, 32), norm.shape); assertTrue(leading.isEmpty()) + assertEquals(rank1.storage, norm.storage) // a view, not a copy + + val rank3 = denseView(Shape(2, 3, 8)) { it.toFloat() } + val (flat, dims) = KernelDispatch.normalizeActivation(rank3) + assertEquals(Shape(6, 8), flat.shape); assertEquals(listOf(2, 3), dims.toList()) + assertEquals(rank3.storage, flat.storage) + assertEquals(rank3.get(1, 2, 3), flat.get(5, 3)) + + val rank2 = denseView(Shape(4, 8)) { it.toFloat() } + assertEquals(rank2, KernelDispatch.normalizeActivation(rank2).first) + } + + @Test + fun theReferenceKernelIsCorrectForARank1ActivationTimesAPackedWeight() { + // the exact #993 repro, through the registry: no special-casing, finite output + val (w, wf) = q8Weight(3) + val x = denseView(Shape(32)) { (it % 5).toFloat() } + val (a, _) = KernelDispatch.normalizeActivation(x) + val out = denseView(Shape(1, 3)) { 0f } + val sink = RecordingTraceSink() + KernelDispatch.matmul(a, w, out, sink = sink) + + for (j in 0 until 3) { + var expect = 0f + for (t in 0 until 32) expect += (t % 5).toFloat() * wf[j * 32 + t] + assertTrue(abs(out.get(0, j) - expect) < 1e-3f, "row $j: ${out.get(0, j)} vs $expect") + assertTrue(out.get(0, j).isFinite()) + } + val run = assertIs(sink.events().single()) + assertEquals("matmul", run.op); assertEquals("reference", run.kernel) + assertEquals("model.layers[0].attn.q_proj.weight", run.inputs[1]!!.canonical) + } + + @Test + fun aRegisteredKernelWinsOverTheReferencePath() { + val (w, _) = q8Weight(2) + val a = denseView(Shape(1, 32)) { 1f } + val out = denseView(Shape(1, 2)) { 0f } + val key = KernelKey.matmul(a, w) + var ran = 0 + KernelDispatch.register(object : ViewKernel { + override val key: KernelKey = key + override val name: String = "fake-q8" + override fun run(inputs: List, out: TensorView) { ran++; out.set(0, 0, value = 7f); out.set(0, 1, value = 8f) } + }) + assertEquals("fake-q8", KernelDispatch.find(key)?.name) + val sink = RecordingTraceSink() + KernelDispatch.matmul(a, w, out, sink = sink) + assertEquals(1, ran); assertEquals(7f, out.get(0, 0)); assertEquals(8f, out.get(0, 1)) + assertEquals("fake-q8", assertIs(sink.events().single()).kernel) + assertNull(KernelDispatch.find(KernelKey("softmax", key.operands))) + } + + @Test + fun aStridedActivationGetsAVisibleGatherAdapter() { + val (w, wf) = q8Weight(1) + val wide = denseView(Shape(2, 64)) { it.toFloat() } + // two rows of 32 taken from a 64-wide buffer: real gaps between rows + val a = wide.narrow(1, 0, 32) + assertTrue(!a.isContiguous, "the activation must be strided for this test to mean anything") + val out = denseView(Shape(2, 1)) { 0f } + val scope = ForwardScope(256) + val sink = RecordingTraceSink() + KernelDispatch.matmul(a, w, out, scope, sink) + val adapter = assertIs(sink.events().first { it is TraceEvent.AdapterInserted }) + assertEquals("gather", adapter.kind); assertEquals(64L * 4, adapter.bytes) // 2 x 32 floats gathered + assertEquals(ScopeKind.FORWARD, scope.kind) + assertTrue(out.get(0, 0).isFinite() && out.get(1, 0).isFinite()) + // the gathered copy is what the kernel read: same numbers as a manual dot product + var expect = 0f + for (t in 0 until 32) expect += a.get(1, t) * wf[t] + assertTrue(abs(out.get(1, 0) - expect) < 1e-3f, "${out.get(1, 0)} vs $expect") + scope.close() + } + + @Test + fun q4kWeightsAlsoGoThroughTheReferencePath() { + // #991's shape: the activation is not the subtype the fast path wanted; the reference decodes anyway + val bytes = ByteArray(144) + for (i in bytes.indices) bytes[i] = (i * 7).toByte() + val d = half(0.02f); bytes[0] = (d and 0xFF).toByte(); bytes[1] = ((d ushr 8) and 0xFF).toByte() + val dmin = half(0.01f); bytes[2] = (dmin and 0xFF).toByte(); bytes[3] = ((dmin ushr 8) and 0xFF).toByte() + val data = Q4_KBlockTensorData(Shape(1, 256), bytes) + val w = TensorView.packed(Storage.Heap.wrap(bytes, mutable = false), Shape(1, 256), TensorEncoding.Q4_K, PackedBlockDecoder(data)) + val a = denseView(Shape(1, 256)) { 0.5f } + val out = denseView(Shape(1, 1)) { 0f } + KernelDispatch.matmul(a, w, out) + var expect = 0f + for (t in 0 until 256) expect += 0.5f * data.toFloatArray()[t] + assertTrue(abs(out.get(0, 0) - expect) < 1e-2f, "${out.get(0, 0)} vs $expect") + } +}