From 59ba3ba5d0a24221d9cb7c5dc9b44633e371bc5c Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 11:05:11 +0200 Subject: [PATCH] =?UTF-8?q?feat(backends):=20exact=20FP32=C3=97b1.58=20mat?= =?UTF-8?q?mul=20dispatch=20via=20the=20ternary=20f32=20pack?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TernaryF32GemvNative is the array-shaped seam over the vendored NeoGPU LUT kernel (#1137), sibling of BitNetGemvNative. TernaryF32KernelPack registers a view kernel under the exact key matmul(FP32 dense × BITNET_B1_58 row-major) — the key KernelDispatch checks before the requantize branch, so installing the pack short-circuits the int8 adapter with zero dispatcher changes, and absence leaves today's behavior untouched (a notice, never a crash; deliberately nothing is registered without the native kernel — the tuned int8 path beats a Kotlin f32 loop as the portable fallback). The view kernel loops the gemv per activation row (prefill included), applies the in-band per-tensor scale the native side never sees, and falls back to the new TernaryF32GemvKernel reference for non-heap storage, strided views, or k % 4 != 0 (the sequential packing crosses byte boundaries between rows then). The reference is pinned against the decoded matmul; the pack contract against a FakeNative; and the real vendored kernel against real dispatch in TernaryF32FfmPackTest — the FFM object now implements the SPI (gemv → gemvPacked) with an install() helper, closing the "FFM has no ternary consumer" gap. Kernel-support matrix intentionally unchanged: pack-registered kernels are not KernelProvider accessors, same as bitnet_gemv. Refs #1138, #1136 Co-Authored-By: Claude Fable 5 --- .../api/kernel/TernaryF32GemvKernel.kt | 93 +++++++++ .../api/kernel/TernaryF32KernelPack.kt | 144 ++++++++++++++ .../api/kernel/TernaryF32GemvKernelTest.kt | 100 ++++++++++ .../api/kernel/TernaryF32KernelPackTest.kt | 182 ++++++++++++++++++ .../exec/kernel/NativeTernaryF32GemvKernel.kt | 23 ++- .../kernel/NativeTernaryF32GemvKernelTest.kt | 18 +- .../exec/kernel/TernaryF32FfmPackTest.kt | 76 ++++++++ 7 files changed, 623 insertions(+), 13 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernel.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPack.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernelTest.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPackTest.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryF32FfmPackTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernel.kt new file mode 100644 index 000000000..34671d15b --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernel.kt @@ -0,0 +1,93 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * The reference `ternary_f32_gemv` (#1138): **FP32 activations** against `BITNET_B1_58` ternary + * weights — exact, no activation quantization. + * + * This is the f32 sibling of [BitNetGemvKernel]. That kernel is the W1.58A8 path: it asks the + * dispatcher to requantize activations to int8-absmax first, trading ~1.5 % quantization error for + * `sdot`-friendly integer math. This one keeps the activations as they are — a ternary weight is + * a `+1`, `-1` or nothing, so the dot product is float adds and subtracts, and the result equals + * the FP32 matmul against the decoded weight bit-for-bit apart from summation order. It serves the + * **exact** dispatch key (FP32 dense × `BITNET_B1_58`), which `KernelDispatch.matmul` checks + * *before* the requantize branch — installing this pack short-circuits the int8 adapter with zero + * dispatcher changes. + * + * Only `BITNET_B1_58` is served: its payload (four codes per byte, low bit-pair first, in element + * order) is what the vendored NeoGPU LUT kernel (#1137) reads, and its scale is in-band. The GGML + * block types (TQ1_0/TQ2_0) have per-block scales and interleaved payloads — they keep their + * int8 path. + * + * Operands: `[rows, k]` FP32 dense activations × `[n, k]` `BITNET_B1_58` weights in canonical + * row-major order. Output `[rows, n]` FP32. The per-tensor scale is applied to the output — + * native implementations ([TernaryF32GemvNative]) do not see it. + */ +@ExperimentalMemoryApi +public class TernaryF32GemvKernel(override val key: KernelKey) : ViewKernel { + + override val name: String get() = "ternary_f32_gemv/reference" + + override fun run(inputs: List, out: TensorView) { + require(inputs.size == 2) { "ternary_f32_gemv takes (activation, weight), got ${inputs.size} operands" } + val a = inputs[0] + val w = inputs[1] + require(a.format.dtype == FP32) { "activation must be FP32, was ${a.format}" } + require(w.format.encoding == TensorEncoding.BITNET_B1_58) { + "weight must be ${TensorEncoding.BITNET_B1_58}, was ${w.format}" + } + require(a.shape.rank == 2 && w.shape.rank == 2 && out.shape.rank == 2) { "ternary_f32_gemv is 2-D" } + val rows = a.shape[0] + val k = a.shape[1] + val n = w.shape[0] + require(w.shape[1] == k) { "inner dimensions differ: activation k=$k, weight k=${w.shape[1]}" } + require(out.shape[0] == rows && out.shape[1] == n) { "out must be [$rows, $n], was ${out.shape}" } + if (rows == 0 || n == 0) return + + val bytes = weightBytes(w) + val byteOffset = (w.storage as Storage.Heap).arrayOffset + // Codes are hoisted out of the row loop, and the per-tensor scale out of everything. + val codes = TernaryCodec.codes(TensorEncoding.BITNET_B1_58, bytes, n * k, byteOffset) + val scale = TernaryCodec.bitNetScale(bytes, n * k, byteOffset) + + for (r in 0 until rows) { + for (o in 0 until n) { + var acc = 0f + val base = o * k + for (i in 0 until k) { + when (codes[base + i].toInt()) { + 1 -> acc += a.get(r, i) + -1 -> acc -= a.get(r, i) + 2 -> acc += 2f * a.get(r, i) // byte code 3; loaders reject it, decode agrees + else -> Unit // zero weights cost nothing + } + } + out.set(r, o, value = acc * scale) + } + } + } + + private fun weightBytes(w: TensorView): ByteArray { + val heap = w.storage as? Storage.Heap + ?: throw UnsupportedOperationException("ternary_f32_gemv reads ternary weights from heap storage in this milestone") + return heap.bytes ?: throw UnsupportedOperationException("ternary weights need byte storage") + } + + public companion object { + /** The exact key this kernel serves: FP32 dense contiguous × `BITNET_B1_58` row-major. */ + public fun keyFor(): KernelKey = KernelKey( + op = "matmul", + operands = listOf( + OperandKey.contiguous(Format.dense(FP32)), + OperandKey(Format(FP32, TensorEncoding.BITNET_B1_58), LayoutClass.BLOCKED_ROW_MAJOR), + ), + ) + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPack.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPack.kt new file mode 100644 index 000000000..6b5935675 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPack.kt @@ -0,0 +1,144 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryCodec + +/** + * The f32-activation ternary gemv a platform pack supplies (#1138) — the seam over the vendored + * NeoGPU LUT kernel (#1137, `skainet_ternary_f32_gemv`). + * + * Array-shaped like [BitNetGemvNative] and for the same reason: implementations are FFM, JNI or + * cinterop shims that pin primitive arrays; [TernaryF32KernelPack] does the view unwrapping once. + * + * The weight is the sequential `BITNET_B1_58` **payload** (four codes per byte, low bit-pair + * first, element order — byte-identical to what [TernaryCodec.encodeBitNet] writes). The + * per-tensor scale is NOT applied here — the wrapping view kernel owns it. `inputDim` must be a + * multiple of 4 (the packing is byte-per-4-elements per row). + */ +@ExperimentalMemoryApi +public interface TernaryF32GemvNative { + /** A name for logs and traces, e.g. `ffm`, `neon`. */ + public val name: String + + /** `out[o] = Σ_k activation(k) · code(o, k)` — one row, scale not applied. */ + public fun gemvPacked( + activation: FloatArray, + activationOffset: Int, + weight: ByteArray, + weightByteOffset: Int, + inputDim: Int, + outputDim: Int, + out: FloatArray, + outOffset: Int, + ) +} + +/** + * Installs the exact FP32 × `BITNET_B1_58` path (#1138). + * + * Registration happens **only when a native kernel is present**, and that is deliberate: without + * the LUT kernel the portable way to multiply a ternary weight is the existing int8-requantize → + * `bitnet_gemv` path ([TernaryKernelPacks]), which stays untouched as the fallback. Registering a + * Kotlin f32 reference into dispatch would *shadow* that tuned path with a slower exact one; the + * f32 reference ([TernaryF32GemvKernel]) exists as the correctness oracle and the in-kernel + * fallback, not as a dispatch entry. + * + * Absence of the native artifact is a notice through [warn], never a crash — behavior is exactly + * today's. + */ +@ExperimentalMemoryApi +public object TernaryF32KernelPack { + + /** What [install] returns when no native kernel is available and nothing was registered. */ + public const val NOT_INSTALLED: String = "ternary_f32_gemv/not-installed" + + /** + * @param native the platform kernel, or `null` when its artifact is absent + * @param capabilities what [native] needs; recorded in the key so a device without them never + * selects it — the vendored kernel itself needs none beyond baseline NEON + * @param warn where the "running without the exact f32 path" notice goes + * @return the name of the kernel that will serve the exact key, or [NOT_INSTALLED] + */ + public fun install( + native: TernaryF32GemvNative? = null, + capabilities: Set = emptySet(), + warn: (String) -> Unit = {}, + ): String { + if (native == null) { + warn( + "ternary_f32_gemv: no native kernel available — FP32×b1.58 matmuls keep the " + + "int8-requantize path. Add the native artifact for the exact f32 path; " + + "nothing else changes.", + ) + return NOT_INSTALLED + } + val kernel = NativeTernaryF32ViewKernel(native, TernaryF32GemvKernel.keyFor().copy(capabilities = capabilities)) + KernelDispatch.register(kernel) + // The dispatcher builds its key from the operands, which say nothing about the CPU — the + // capability-free key is the reachable one, the capability key documents the requirement. + // Same two-key pattern as TernaryKernelPacks.install. + KernelDispatch.register(NativeTernaryF32ViewKernel(native, TernaryF32GemvKernel.keyFor())) + return kernel.name + } +} + +/** + * A [ViewKernel] over a [TernaryF32GemvNative]: unwraps the views once, loops the native gemv + * over the activation rows (prefill included — each row is an independent gemv), and applies the + * `BITNET_B1_58` per-tensor scale to what the native kernel wrote. + * + * Falls back to the reference for anything the native contract does not cover — non-heap storage, + * a strided view, or `k % 4 != 0` (the sequential packing crosses byte boundaries between rows + * then) — instead of failing: the fast path is an optimization, never a correctness requirement. + */ +@ExperimentalMemoryApi +public class NativeTernaryF32ViewKernel( + private val native: TernaryF32GemvNative, + override val key: KernelKey, +) : ViewKernel { + + override val name: String get() = "ternary_f32_gemv/${native.name}" + + private val reference = TernaryF32GemvKernel(key) + + override fun run(inputs: List, out: TensorView) { + val a = inputs[0] + val w = inputs[1] + val rows = a.shape[0] + val k = a.shape[1] + val n = w.shape[0] + val activationFloats = (a.storage as? Storage.Heap)?.floats + val weightBytes = (w.storage as? Storage.Heap)?.bytes + val outFloats = (out.storage as? Storage.Heap)?.floats + if (k % 4 != 0 || activationFloats == null || weightBytes == null || outFloats == null || + !a.isContiguous || !out.isContiguous + ) { + reference.run(inputs, out) + return + } + if (rows == 0 || n == 0) return + val aOffset = (a.storage as Storage.Heap).arrayOffset + val wOffset = (w.storage as Storage.Heap).arrayOffset + val outOffset = (out.storage as Storage.Heap).arrayOffset + for (r in 0 until rows) { + native.gemvPacked( + activation = activationFloats, + activationOffset = aOffset + r * k, + weight = weightBytes, + weightByteOffset = wOffset, + inputDim = k, + outputDim = n, + out = outFloats, + outOffset = outOffset + r * n, + ) + } + // The native kernel computes the unscaled codes-dot; the per-tensor scale lives in the + // weight's trailing FP32 and is applied once, here. + val scale = TernaryCodec.bitNetScale(weightBytes, n * k, wOffset) + if (scale != 1f) { + for (i in outOffset until outOffset + rows * n) outFloats[i] *= scale + } + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernelTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernelTest.kt new file mode 100644 index 000000000..e50f06d0b --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32GemvKernelTest.kt @@ -0,0 +1,100 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryBlockDecoder +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1138: the f32 reference against the one thing it must equal — an FP32 matmul over the + * *decoded* weight ([TernaryCodec.decodeBitNet]). The codes-dot is exact; only summation + * order may differ. + */ +@OptIn(ExperimentalMemoryApi::class) +class TernaryF32GemvKernelTest { + + private fun ternaryValues(count: Int, seed: Int): FloatArray { + var s = seed + return FloatArray(count) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 3 - 1) * 0.7f + } + } + + private fun weight(n: Int, k: Int, seed: Int = 5): TensorView { + val bytes = TernaryCodec.encodeBitNet(ternaryValues(n * k, seed)) + return TensorView.packed( + Storage.Heap.wrap(bytes), Shape(n, k), TensorEncoding.BITNET_B1_58, + TernaryBlockDecoder(TensorEncoding.BITNET_B1_58, n * k), + ) + } + + private fun activation(rows: Int, k: Int, seed: Int = 9): TensorView { + var s = seed + val floats = FloatArray(rows * k) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2000 - 1000) / 1000f + } + return TensorView.dense(Storage.Heap.wrap(floats), Shape(rows, k), FP32) + } + + @Test + fun referenceEqualsMatmulOverTheDecodedWeight() { + val rows = 2; val k = 96; val n = 5 + val w = weight(n, k) + val a = activation(rows, k) + val out = TensorView.dense(Storage.Heap.floats(rows * n), Shape(rows, n), FP32) + TernaryF32GemvKernel(TernaryF32GemvKernel.keyFor()).run(listOf(a, w), out) + + val bytes = (w.storage as Storage.Heap).bytes!! + val decoded = TernaryCodec.decodeBitNet(bytes, n * k) + for (r in 0 until rows) { + for (o in 0 until n) { + var want = 0f + for (i in 0 until k) want += a.get(r, i) * decoded[o * k + i] + val got = out.get(r, o) + assertTrue( + abs(got - want) <= 1e-4f * maxOf(1f, abs(want)), + "[$r,$o]: reference=$got decoded-matmul=$want", + ) + } + } + } + + @Test + fun kIndivisibleByFourStillMatchesTheDecodedMatmul() { + // BITNET_B1_58 packs the flattened tensor, so k % 4 != 0 crosses byte + // boundaries between rows — the reference reads linear codes and must + // not care. (The native view kernel falls back to this path.) + val rows = 1; val k = 6; val n = 3 + val w = weight(n, k, seed = 11) + val a = activation(rows, k, seed = 13) + val out = TensorView.dense(Storage.Heap.floats(rows * n), Shape(rows, n), FP32) + TernaryF32GemvKernel(TernaryF32GemvKernel.keyFor()).run(listOf(a, w), out) + + val decoded = TernaryCodec.decodeBitNet((w.storage as Storage.Heap).bytes!!, n * k) + for (o in 0 until n) { + var want = 0f + for (i in 0 until k) want += a.get(0, i) * decoded[o * k + i] + assertTrue(abs(out.get(0, o) - want) <= 1e-5f, "[$o]: ${out.get(0, o)} vs $want") + } + } + + @Test + fun mismatchedInnerDimensionsAreRejected() { + val w = weight(n = 2, k = 8) + val a = activation(rows = 1, k = 12) + val out = TensorView.dense(Storage.Heap.floats(2), Shape(1, 2), FP32) + assertFailsWith { + TernaryF32GemvKernel(TernaryF32GemvKernel.keyFor()).run(listOf(a, w), out) + } + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPackTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPackTest.kt new file mode 100644 index 000000000..d302655c7 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryF32KernelPackTest.kt @@ -0,0 +1,182 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryBlockDecoder +import sk.ainet.lang.memory.TernaryCodec +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.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1138: the exact FP32×`BITNET_B1_58` path takes over dispatch when its native kernel is present, + * and its absence changes nothing — FP32 activations keep flowing through the int8-requantize → + * `bitnet_gemv` path exactly as before, with a notice instead of a crash. + */ +@OptIn(ExperimentalMemoryApi::class) +class TernaryF32KernelPackTest { + + private val k = 64 + private val n = 4 + + @BeforeTest fun setUp() = KernelDispatch.clearForTesting() + @AfterTest fun tearDown() = KernelDispatch.clearForTesting() + + /** A stand-in for the FFM/JNI kernel: records calls, computes the unscaled codes-dot. */ + private class FakeNative(override val name: String = "fake-lut") : TernaryF32GemvNative { + var calls: Int = 0 + override fun gemvPacked( + activation: FloatArray, activationOffset: Int, + weight: ByteArray, weightByteOffset: Int, + inputDim: Int, outputDim: Int, + out: FloatArray, outOffset: Int, + ) { + calls++ + for (o in 0 until outputDim) { + var acc = 0f + for (i in 0 until inputDim) { + val element = o * inputDim + i + val code = ((weight[weightByteOffset + element / 4].toInt() + shr ((element % 4) * 2)) and 3) - 1 + acc += code * activation[activationOffset + i] + } + out[outOffset + o] = acc + } + } + } + + private fun weight(): TensorView { + var seed = 5 + val values = FloatArray(n * k) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 3 - 1) * 0.5f + } + val bytes = TernaryCodec.encodeBitNet(values) + return TensorView.packed( + Storage.Heap.wrap(bytes), Shape(n, k), TensorEncoding.BITNET_B1_58, + TernaryBlockDecoder(TensorEncoding.BITNET_B1_58, n * k), + ) + } + + private fun activation(rows: Int = 1): TensorView { + var seed = 9 + val floats = FloatArray(rows * k) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 2000 - 1000) / 1000f + } + return TensorView.dense(Storage.Heap.wrap(floats), Shape(rows, k), FP32) + } + + @Test + fun withoutTheArtifactTheInt8PathServesUnchangedAndTheCallerIsTold() { + TernaryKernelPacks.install(native = null, warn = {}) + + val warnings = mutableListOf() + val serving = TernaryF32KernelPack.install(native = null, warn = { warnings += it }) + assertEquals(TernaryF32KernelPack.NOT_INSTALLED, serving) + assertEquals(1, warnings.size, "exactly one notice, not a crash: $warnings") + assertTrue(warnings.single().contains("int8-requantize"), warnings.single()) + + // FP32 × b1.58 dispatch behaves exactly as before this pack existed: + // requantize adapter + the int8 reference. + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + val sink = RecordingTraceSink() + KernelDispatch.matmul(activation(), weight(), out, Scope.Ambient, sink) + assertEquals("bitnet_gemv/reference", sink.eventsOf().single().kernel) + } + + @Test + fun withTheArtifactTheExactKeyBeatsTheRequantizePath() { + TernaryKernelPacks.install(native = null, warn = {}) + val native = FakeNative() + val serving = TernaryF32KernelPack.install(native) + assertEquals("ternary_f32_gemv/fake-lut", serving) + + val w = weight() + val a = activation() + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + val sink = RecordingTraceSink() + KernelDispatch.matmul(a, w, out, Scope.Ambient, sink) + + assertEquals("ternary_f32_gemv/fake-lut", sink.eventsOf().single().kernel) + assertEquals(1, native.calls) + assertTrue( + sink.eventsOf().isEmpty(), + "the exact f32 path needs no requantize adapter", + ) + + // and the scaled result equals the f32 reference + val fromReference = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + TernaryF32GemvKernel(TernaryF32GemvKernel.keyFor()).run(listOf(a, w), fromReference) + for (o in 0 until n) { + val got = out.get(0, o) + val want = fromReference.get(0, o) + assertTrue(abs(got - want) <= 1e-5f * maxOf(1f, abs(want)), "[$o]: $got vs $want") + } + } + + @Test + fun multiRowActivationsLoopTheNativeGemvPerRow() { + val native = FakeNative() + TernaryF32KernelPack.install(native) + val rows = 3 + val w = weight() + val a = activation(rows) + val out = TensorView.dense(Storage.Heap.floats(rows * n), Shape(rows, n), FP32) + NativeTernaryF32ViewKernel(native, TernaryF32GemvKernel.keyFor()).run(listOf(a, w), out) + assertEquals(rows, native.calls, "prefill loops the gemv, one call per row") + + val fromReference = TensorView.dense(Storage.Heap.floats(rows * n), Shape(rows, n), FP32) + TernaryF32GemvKernel(TernaryF32GemvKernel.keyFor()).run(listOf(a, w), fromReference) + for (r in 0 until rows) for (o in 0 until n) { + val got = out.get(r, o) + val want = fromReference.get(r, o) + assertTrue(abs(got - want) <= 1e-5f * maxOf(1f, abs(want)), "[$r,$o]: $got vs $want") + } + } + + @Test + fun kIndivisibleByFourFallsBackToTheReferenceInsteadOfFailing() { + val native = FakeNative() + val oddK = 6 + var seed = 3 + val values = FloatArray(n * oddK) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 3 - 1).toFloat() + } + val w = TensorView.packed( + Storage.Heap.wrap(TernaryCodec.encodeBitNet(values)), Shape(n, oddK), + TensorEncoding.BITNET_B1_58, TernaryBlockDecoder(TensorEncoding.BITNET_B1_58, n * oddK), + ) + val floats = FloatArray(oddK) { (it + 1).toFloat() } + val a = TensorView.dense(Storage.Heap.wrap(floats), Shape(1, oddK), FP32) + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + NativeTernaryF32ViewKernel(native, TernaryF32GemvKernel.keyFor()).run(listOf(a, w), out) + assertEquals(0, native.calls, "the packing crosses byte boundaries — reference serves") + + val decoded = TernaryCodec.decodeBitNet((w.storage as Storage.Heap).bytes!!, n * oddK) + for (o in 0 until n) { + var want = 0f + for (i in 0 until oddK) want += floats[i] * decoded[o * oddK + i] + assertTrue(abs(out.get(0, o) - want) <= 1e-5f, "[$o]: ${out.get(0, o)} vs $want") + } + } + + @Test + fun theCapabilityIsRecordedInTheKey() { + TernaryF32KernelPack.install(FakeNative(), setOf("ffm")) + val keys = KernelDispatch.kernels().filter { it.name.startsWith("ternary_f32_gemv/fake") }.map { it.key } + assertTrue(keys.any { it.capabilities == setOf("ffm") }, "the pack declares what it needs: $keys") + assertTrue(keys.any { it.capabilities.isEmpty() }, "and is reachable from an operand-only key: $keys") + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernel.kt index 89b142c69..2ab7d0bef 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernel.kt @@ -6,6 +6,9 @@ import java.lang.foreign.Linker import java.lang.foreign.MemorySegment import java.lang.foreign.ValueLayout import java.lang.invoke.MethodHandle +import sk.ainet.backend.api.kernel.TernaryF32GemvNative +import sk.ainet.backend.api.kernel.TernaryF32KernelPack +import sk.ainet.lang.memory.ExperimentalMemoryApi /** * Native (FFM) downcall to the vendored NeoGPU ternary LUT kernel. @@ -30,14 +33,26 @@ import java.lang.invoke.MethodHandle * pinned by `NativeTernaryF32GemvKernelTest`. * * Refs SKaiNET issue #1137 (vendored from anjaustin/neogpu, MIT — see - * native/src/vendor/neogpu/README.md). The `TernaryF32GemvNative` SPI - * wiring into KernelDispatch follows in #1138. + * native/src/vendor/neogpu/README.md); implements the [TernaryF32GemvNative] + * seam so [install] can hand it to `TernaryF32KernelPack` (#1138) — the + * first ternary FFM consumer. */ -internal object NativeTernaryF32GemvKernel { +@OptIn(ExperimentalMemoryApi::class) +internal object NativeTernaryF32GemvKernel : TernaryF32GemvNative { + + override val name: String get() = "ffm" fun isAvailable(): Boolean = handle != null - fun gemv( + /** + * Register this kernel with [TernaryF32KernelPack] when the bundled + * library resolves; without it the pack warns and dispatch keeps the + * int8-requantize path. Returns the serving kernel name. + */ + fun install(warn: (String) -> Unit = {}): String = + TernaryF32KernelPack.install(if (isAvailable()) this else null, warn = warn) + + override fun gemvPacked( input: FloatArray, inputOffset: Int, weight: ByteArray, weightByteOffset: Int, inputDim: Int, outputDim: Int, diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernelTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernelTest.kt index bda76a9cc..f4a46f68c 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernelTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeTernaryF32GemvKernelTest.kt @@ -63,7 +63,7 @@ class NativeTernaryF32GemvKernelTest { val weight = ByteArray(256) { it.toByte() } val expected = referenceGemv(input, weight, 0, inputDim, 1) val out = FloatArray(1) - NativeTernaryF32GemvKernel.gemv(input, 0, weight, 0, inputDim, 1, out, 0) + NativeTernaryF32GemvKernel.gemvPacked(input, 0, weight, 0, inputDim, 1, out, 0) assertEquals(expected[0], out[0], "all-bytes golden must match bit-exactly") } @@ -71,7 +71,7 @@ class NativeTernaryF32GemvKernelTest { fun byte_code_3_decodes_to_plus_two() { // 0xFF = four 2-bit codes of 3 → each lane decodes to +2.0. val out = FloatArray(1) - NativeTernaryF32GemvKernel.gemv( + NativeTernaryF32GemvKernel.gemvPacked( floatArrayOf(1f, 1f, 1f, 1f), 0, byteArrayOf(0xFF.toByte()), 0, 4, 1, out, 0, @@ -88,7 +88,7 @@ class NativeTernaryF32GemvKernelTest { val weight = ByteArray(outputDim * inputDim / 4).also { rng.nextBytes(it) } val expected = referenceGemv(input, weight, 0, inputDim, outputDim) val out = FloatArray(outputDim) - NativeTernaryF32GemvKernel.gemv(input, 0, weight, 0, inputDim, outputDim, out, 0) + NativeTernaryF32GemvKernel.gemvPacked(input, 0, weight, 0, inputDim, outputDim, out, 0) for (i in out.indices) { val diff = abs(expected[i] - out[i]) assertTrue( @@ -112,11 +112,11 @@ class NativeTernaryF32GemvKernelTest { val weight = ByteArray(outputDim * rowBytes).also { rng.nextBytes(it) } val threaded = FloatArray(outputDim) - NativeTernaryF32GemvKernel.gemv(input, 0, weight, 0, inputDim, outputDim, threaded, 0) + NativeTernaryF32GemvKernel.gemvPacked(input, 0, weight, 0, inputDim, outputDim, threaded, 0) val perRow = FloatArray(outputDim) for (n in 0 until outputDim) { - NativeTernaryF32GemvKernel.gemv( + NativeTernaryF32GemvKernel.gemvPacked( input, 0, weight, n * rowBytes, inputDim, 1, perRow, n, ) } @@ -139,7 +139,7 @@ class NativeTernaryF32GemvKernelTest { weight[5] = 0x22 weight[6] = 0x22 val out = FloatArray(4) { -1f } - NativeTernaryF32GemvKernel.gemv(input, pad, weight, 5, inputDim, 2, out, 2) + NativeTernaryF32GemvKernel.gemvPacked(input, pad, weight, 5, inputDim, 2, out, 2) // in (after offset) = 1..8; row0 decode = {+1,-1,+1,-1, +1,-1,+1,-1} // → 1-2+3-4+5-6+7-8 = -4; row1 all zeros → 0. assertEquals(-1f, out[0]); assertEquals(-1f, out[1]) @@ -149,7 +149,7 @@ class NativeTernaryF32GemvKernelTest { @Test fun rejects_non_multiple_of_4_input_dim() { assertFailsWith { - NativeTernaryF32GemvKernel.gemv( + NativeTernaryF32GemvKernel.gemvPacked( FloatArray(6), 0, ByteArray(2), 0, 6, 1, FloatArray(1), 0, ) } @@ -157,7 +157,7 @@ class NativeTernaryF32GemvKernelTest { @Test fun zero_output_dim_is_no_op() { - NativeTernaryF32GemvKernel.gemv( + NativeTernaryF32GemvKernel.gemvPacked( FloatArray(4) { 1f }, 0, ByteArray(1), 0, 4, 0, FloatArray(0), 0, ) } @@ -165,7 +165,7 @@ class NativeTernaryF32GemvKernelTest { @Test fun zero_input_dim_zeros_output() { val out = FloatArray(3) { 9f } - NativeTernaryF32GemvKernel.gemv( + NativeTernaryF32GemvKernel.gemvPacked( FloatArray(0), 0, ByteArray(0), 0, 0, 3, out, 0, ) for (v in out) assertEquals(0f, v, "output should be zeroed for inputDim=0") diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryF32FfmPackTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryF32FfmPackTest.kt new file mode 100644 index 000000000..7ba30139a --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryF32FfmPackTest.kt @@ -0,0 +1,76 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.backend.api.kernel.TernaryF32GemvKernel +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Scope +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryBlockDecoder +import sk.ainet.lang.memory.TernaryCodec +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.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1138 end-to-end on the JVM: the REAL vendored kernel behind the REAL dispatcher. Install the + * FFM pack, hand `KernelDispatch.matmul` an FP32 activation and a `BITNET_B1_58` weight view, and + * the NeoGPU LUT kernel serves the exact key — no requantize adapter, results equal to the f32 + * reference. + */ +@OptIn(ExperimentalMemoryApi::class) +class TernaryF32FfmPackTest { + + private val k = 2560 // BitNet-2B hidden size + private val n = 8 + + @BeforeTest fun setUp() { + KernelDispatch.clearForTesting() + assertTrue(NativeTernaryF32GemvKernel.isAvailable(), "bundled libskainet_kernels must resolve") + } + @AfterTest fun tearDown() = KernelDispatch.clearForTesting() + + @Test + fun theVendoredKernelServesDispatchAndMatchesTheReference() { + val serving = NativeTernaryF32GemvKernel.install() + assertEquals("ternary_f32_gemv/ffm", serving) + + var seed = 5 + val values = FloatArray(n * k) { + seed = seed * 1103515245 + 12345 + ((seed ushr 16) % 3 - 1) * 0.5f + } + val w = TensorView.packed( + Storage.Heap.wrap(TernaryCodec.encodeBitNet(values)), Shape(n, k), + TensorEncoding.BITNET_B1_58, TernaryBlockDecoder(TensorEncoding.BITNET_B1_58, n * k), + ) + seed = 9 + val a = TensorView.dense( + Storage.Heap.wrap(FloatArray(k) { seed = seed * 1103515245 + 12345; ((seed ushr 16) % 2000 - 1000) / 1000f }), + Shape(1, k), FP32, + ) + + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + val sink = RecordingTraceSink() + KernelDispatch.matmul(a, w, out, Scope.Ambient, sink) + + assertEquals("ternary_f32_gemv/ffm", sink.eventsOf().single().kernel) + assertTrue(sink.eventsOf().isEmpty(), "no requantize adapter on the exact path") + + val fromReference = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + TernaryF32GemvKernel(TernaryF32GemvKernel.keyFor()).run(listOf(a, w), fromReference) + for (o in 0 until n) { + val got = out.get(0, o) + val want = fromReference.get(0, o) + assertTrue(abs(got - want) <= 1e-3f * maxOf(1f, abs(want)), "[$o]: ffm=$got reference=$want") + } + } +}