From 9d2664e408be5d63a5aa1b101a5278ea3537b4dd Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 15:34:41 +0200 Subject: [PATCH] =?UTF-8?q?feat(lang,backends,io):=20BITNET=5FPLANES=20?= =?UTF-8?q?=E2=80=94=20the=20multi-plane=20lm=5Fhead=20format?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TensorEncoding.BITNET_PLANES is NeoGPU's lm_head weight format as a plain SKaiNET encoding (#1150): 8 sequentially-packed trit planes + one FP16 scale per row, plane p worth 1/3^p — "16 bits as eight ternary digits", truncation error <= rowScale/(2*3^7). The point is speed, not memory (2 B/weight, FP16-sized): the vendored fused 4-plane LUT kernel reads planes 0-3 in one baseline-NEON pass, and an application can rescore top candidates with planes 4-7 (NeoGPU's two-stage lm_head) through the codec's row accessors. The format stays format-driven end to end, per the #1136 physiology: no new op, layer, or fusion pass anywhere. - TernaryCodec.encodeBitNetPlanes is the Kotlin port of NeoGPU's hs_mlt_lmhead_encode — per-row absmax as FP16 (decomposed against the FP16-rounded scale the decoder will read, not the exact one), repeated round-to-trit with x3 residual; decodeBitNetPlanes/Row + planesRowScale are the reference readers. BitNetPlanesTensorData blocks per ROW. - BITNET_PLANES deliberately has no BlockSpec and no int8 activation hint: its geometry is row-scoped (row-count-free BlockSpec cannot say it), and without the pack dispatch must fall to the decoding reference, never the requantize adapter — pinned by test. - TernaryLmheadNative + TernaryPlanesKernelPack: the native seam computes 4 fused planes per call (skainet_ternary_lmhead_stage1, exported in #1137); the view kernel makes two calls combined as s0 + s4/81, so dispatch's matmul equals the decoded 8-plane matmul EXACTLY — stage-1 truncation is an application decision, never a dispatch surprise. - FFM face NativeTernaryLmheadKernel (row scales = the weight segment sliced at the aligned scale offset); JNI/K-N faces are follow-ups within #1150. - StreamingGgufParametersLoader gains its first requantizer: RequantizeTo(BITNET_PLANES) from F32/F16/BF16/quantized/I2_S sources, OUT_IN orientation required (the scales are per output row), traced as requantize-planes. Every other RequantizeTo still fails eagerly. Tests: codec truncation-bound goldens + plane-0 sign structure; pack contract with a FakeNative (two-call exactness vs the 8-plane reference, nothing registered without native, no requantize adapter without the pack); the REAL vendored kernel behind REAL dispatch at BitNet hidden size; loader requantize round-trip within the bound + eager rejections. Refs #1150, #1136 Co-Authored-By: Claude Fable 5 --- .../api/kernel/TernaryPlanesKernelPack.kt | 190 ++++++++++++++++++ .../api/kernel/TernaryPlanesKernelPackTest.kt | 143 +++++++++++++ .../exec/kernel/NativeTernaryLmheadKernel.kt | 112 +++++++++++ .../ainet/exec/kernel/TernaryPlanesFfmTest.kt | 67 ++++++ .../io/gguf/StreamingGgufParametersLoader.kt | 96 ++++++++- .../ainet/io/gguf/PlanesRequantizeLoadTest.kt | 106 ++++++++++ .../kotlin/sk/ainet/lang/memory/BlockSpec.kt | 6 + .../sk/ainet/lang/memory/TernaryCodec.kt | 84 ++++++++ .../tensor/data/BitNetPlanesTensorData.kt | 82 ++++++++ .../lang/tensor/storage/TensorEncoding.kt | 40 ++++ .../lang/tensor/data/BitNetPlanesCodecTest.kt | 86 ++++++++ 11 files changed, 1002 insertions(+), 10 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPack.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPackTest.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryLmheadKernel.kt create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryPlanesFfmTest.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PlanesRequantizeLoadTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesTensorData.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesCodecTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPack.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPack.kt new file mode 100644 index 000000000..d2ef09a8e --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPack.kt @@ -0,0 +1,190 @@ +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 fused multi-plane lm_head gemv a platform pack supplies (#1150) — the seam over the + * vendored NeoGPU `skainet_ternary_lmhead_stage1` (#1137). Array-shaped like the other ternary + * seams; the pack does the view unwrapping once. + * + * One call computes **planes [firstPlane, firstPlane+4)** of a `BITNET_PLANES` weight with the + * fused plane weights `{1, ⅓, ⅑, 1/27}` and the FP16 row scales applied: + * + * `out[o] = rowScale[o] · Σ_{q=0}^{3} (1/3^q) · Σ_k in[k] · code(firstPlane+q, o, k)` + * + * The full 8-plane matmul is two calls (`firstPlane` 0 and 4) combined as `s0 + s4 / 81` — which + * is how [NativeTernaryPlanesViewKernel] keeps the dispatch invariant *matmul == decoded matmul*. + * `inputDim % 4 == 0`; [rowScaleByteOffset] must be 2-byte aligned. + */ +@ExperimentalMemoryApi +public interface TernaryLmheadNative { + /** A name for logs and traces, e.g. `ffm`. */ + public val name: String + + public fun lmheadStage1( + activation: FloatArray, + activationOffset: Int, + weight: ByteArray, + planesByteOffset: Int, + planeStrideBytes: Int, + rowScaleByteOffset: Int, + inputDim: Int, + outputDim: Int, + out: FloatArray, + outOffset: Int, + ) +} + +/** + * Installs the exact FP32 × `BITNET_PLANES` path (#1150). Same contract as + * [TernaryF32KernelPack]: registration **only with a native kernel** — without it, dispatch falls + * back to the decoding reference matmul (correct, slow), told through [warn], never a crash. + */ +@ExperimentalMemoryApi +public object TernaryPlanesKernelPack { + + /** What [install] returns when no native kernel is available and nothing was registered. */ + public const val NOT_INSTALLED: String = "ternary_planes_matmul/not-installed" + + public fun install( + native: TernaryLmheadNative? = null, + capabilities: Set = emptySet(), + warn: (String) -> Unit = {}, + ): String { + if (native == null) { + warn( + "ternary_planes_matmul: no native kernel available — BITNET_PLANES matmuls decode " + + "through the reference path. Add the native artifact for the fused LUT path; " + + "nothing else changes.", + ) + return NOT_INSTALLED + } + val kernel = NativeTernaryPlanesViewKernel(native, TernaryPlanesMatmulKernel.keyFor().copy(capabilities = capabilities)) + KernelDispatch.register(kernel) + KernelDispatch.register(NativeTernaryPlanesViewKernel(native, TernaryPlanesMatmulKernel.keyFor())) + return kernel.name + } +} + +/** + * The portable reference for FP32 × `BITNET_PLANES` — decodes each weight row (all 8 planes × + * row scale) through [TernaryCodec] and accumulates in FP32. The correctness oracle and the + * in-kernel fallback; deliberately not a dispatch entry on its own. + */ +@ExperimentalMemoryApi +public class TernaryPlanesMatmulKernel(override val key: KernelKey) : ViewKernel { + + override val name: String get() = "ternary_planes_matmul/reference" + + override fun run(inputs: List, out: TensorView) { + require(inputs.size == 2) { "ternary_planes_matmul takes (activation, weight), got ${inputs.size}" } + val a = inputs[0] + val w = inputs[1] + require(w.format.encoding == TensorEncoding.BITNET_PLANES) { + "weight must be ${TensorEncoding.BITNET_PLANES}, was ${w.format}" + } + require(a.shape.rank == 2 && w.shape.rank == 2 && out.shape.rank == 2) { "ternary_planes_matmul 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 heap = w.storage as? Storage.Heap + ?: throw UnsupportedOperationException("ternary_planes_matmul reads weights from heap storage in this milestone") + val bytes = heap.bytes ?: throw UnsupportedOperationException("BITNET_PLANES weights need byte storage") + val byteOffset = heap.arrayOffset + + val decodedRow = FloatArray(k) + for (o in 0 until n) { + TernaryCodec.decodeBitNetPlanesRow(bytes, n, k, o, decodedRow, 0, byteOffset) + for (r in 0 until rows) { + var acc = 0f + for (i in 0 until k) acc += a.get(r, i) * decodedRow[i] + out.set(r, o, value = acc) + } + } + } + + public companion object { + /** The exact key: FP32 dense contiguous × `BITNET_PLANES` row-major. */ + public fun keyFor(): KernelKey = KernelKey( + op = "matmul", + operands = listOf( + OperandKey.contiguous(Format.dense(FP32)), + OperandKey(Format(FP32, TensorEncoding.BITNET_PLANES), LayoutClass.BLOCKED_ROW_MAJOR), + ), + ) + } +} + +/** + * A [ViewKernel] over a [TernaryLmheadNative]: two fused stage-1 calls per activation row — + * planes 0–3 and planes 4–7 — combined as `s0 + s4 / 81`, so the result is the **full 8-plane** + * matmul and the dispatch invariant *matmul == decoded matmul* holds exactly. (Stage-1-only + * scoring with top-k rescoring is an application-level decision made through the codec's row + * accessors, never through dispatch.) + * + * Falls back to the reference for non-heap storage, strided views, or `k % 4 != 0`. + */ +@ExperimentalMemoryApi +public class NativeTernaryPlanesViewKernel( + private val native: TernaryLmheadNative, + override val key: KernelKey, +) : ViewKernel { + + override val name: String get() = "ternary_planes_matmul/${native.name}" + + private val reference = TernaryPlanesMatmulKernel(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 + val planeStride = TensorEncoding.BITNET_PLANES.planeStrideBytes(n, k) + val scalesOffset = TensorEncoding.BITNET_PLANES.rowScalesByteOffset(n, k) + val high = FloatArray(n) + for (r in 0 until rows) { + // planes 0–3, row scales applied by the kernel + native.lmheadStage1( + activation = activationFloats, activationOffset = aOffset + r * k, + weight = weightBytes, planesByteOffset = wOffset, + planeStrideBytes = planeStride, rowScaleByteOffset = wOffset + scalesOffset, + inputDim = k, outputDim = n, + out = outFloats, outOffset = outOffset + r * n, + ) + // planes 4–7, same fused weights — worth 1/3⁴ of the total + native.lmheadStage1( + activation = activationFloats, activationOffset = aOffset + r * k, + weight = weightBytes, planesByteOffset = wOffset + planeStride * 4, + planeStrideBytes = planeStride, rowScaleByteOffset = wOffset + scalesOffset, + inputDim = k, outputDim = n, + out = high, outOffset = 0, + ) + val base = outOffset + r * n + for (o in 0 until n) outFloats[base + o] += high[o] / 81f + } + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPackTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPackTest.kt new file mode 100644 index 000000000..108cb8ae2 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/TernaryPlanesKernelPackTest.kt @@ -0,0 +1,143 @@ +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.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.data.BitNetPlanesTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1150: the `BITNET_PLANES` pack contract. The native seam computes 4 fused planes per call; + * the view kernel's two-call combination (`s0 + s4/81`) must equal the full 8-plane reference — + * the dispatch invariant *matmul == decoded matmul* holds exactly. + */ +@OptIn(ExperimentalMemoryApi::class) +class TernaryPlanesKernelPackTest { + + private val k = 32 + private val n = 5 + + @BeforeTest fun setUp() = KernelDispatch.clearForTesting() + @AfterTest fun tearDown() = KernelDispatch.clearForTesting() + + /** Computes exactly the 4-plane fused contract, straight from the buffer layout. */ + private class FakeNative(override val name: String = "fake") : TernaryLmheadNative { + var calls: Int = 0 + override fun lmheadStage1( + activation: FloatArray, activationOffset: Int, + weight: ByteArray, planesByteOffset: Int, + planeStrideBytes: Int, rowScaleByteOffset: Int, + inputDim: Int, outputDim: Int, + out: FloatArray, outOffset: Int, + ) { + calls++ + val rowBytes = inputDim / 4 + for (o in 0 until outputDim) { + val scaleBits = (weight[rowScaleByteOffset + o * 2].toInt() and 0xFF) or + ((weight[rowScaleByteOffset + o * 2 + 1].toInt() and 0xFF) shl 8) + val scale = sk.ainet.lang.types.Fp16Codec.decode(scaleBits) + var acc = 0f + var w = 1f + for (q in 0 until 4) { + val base = planesByteOffset + q * planeStrideBytes + o * rowBytes + var dot = 0f + for (i in 0 until inputDim) { + val code = ((weight[base + i / 4].toInt() and 0xFF) shr ((i % 4) * 2)) and 3 + dot += (code - 1) * activation[activationOffset + i] + } + acc += dot * w + w /= 3f + } + out[outOffset + o] = acc * scale + } + } + } + + private fun weight(): TensorView { + val rng = Random(5) + val values = FloatArray(n * k) { (rng.nextFloat() - 0.5f) * 2f } + return BitNetPlanesTensorData.fromFloats(Shape(n, k), values).packedView + } + + private fun activation(rows: Int = 1): TensorView { + val rng = Random(9) + return TensorView.dense( + Storage.Heap.wrap(FloatArray(rows * k) { rng.nextFloat() - 0.5f }), + Shape(rows, k), FP32, + ) + } + + @Test + fun withTheArtifactTheExactKeyServesAndEqualsTheFullPlaneReference() { + val native = FakeNative() + val serving = TernaryPlanesKernelPack.install(native) + assertEquals("ternary_planes_matmul/fake", 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_planes_matmul/fake", sink.eventsOf().single().kernel) + assertEquals(2, native.calls, "full result = planes 0–3 + planes 4–7 / 81 — two fused calls") + + val fromReference = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + TernaryPlanesMatmulKernel(TernaryPlanesMatmulKernel.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-4f * maxOf(1f, abs(want)), "[$o]: $got vs $want") + } + } + + @Test + fun withoutTheArtifactNothingIsRegisteredAndTheCallerIsTold() { + val warnings = mutableListOf() + val serving = TernaryPlanesKernelPack.install(native = null, warn = { warnings += it }) + assertEquals(TernaryPlanesKernelPack.NOT_INSTALLED, serving) + assertEquals(1, warnings.size) + assertTrue(KernelDispatch.kernels().isEmpty(), "nothing registered without the native kernel") + } + + @Test + fun theReferenceEqualsTheDecodedMatmul() { + val w = weight() + val a = activation(rows = 2) + val out = TensorView.dense(Storage.Heap.floats(2 * n), Shape(2, n), FP32) + TernaryPlanesMatmulKernel(TernaryPlanesMatmulKernel.keyFor()).run(listOf(a, w), out) + + val bytes = (w.storage as Storage.Heap).bytes!! + val decoded = TernaryCodec.decodeBitNetPlanes(bytes, n, k) + for (r in 0 until 2) for (o in 0 until n) { + var want = 0f + for (i in 0 until k) want += a.get(r, i) * decoded[o * k + i] + assertTrue(abs(out.get(r, o) - want) <= 1e-5f, "[$r,$o]: ${out.get(r, o)} vs $want") + } + } + + @Test + fun encodingWithoutBlockSpecTriggersNoRequantizeAdapter() { + // BITNET_PLANES deliberately declares no activation hint: without the pack the dispatcher + // must fall to the decoding reference matmul, never the int8 requantize path. + 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) + val kernel = sink.eventsOf().single().kernel + assertEquals("reference", kernel, "decoding reference serves without the pack") + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryLmheadKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryLmheadKernel.kt new file mode 100644 index 000000000..ee6a610e3 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeTernaryLmheadKernel.kt @@ -0,0 +1,112 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.FunctionDescriptor +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.TernaryLmheadNative +import sk.ainet.backend.api.kernel.TernaryPlanesKernelPack +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * Native (FFM) downcall to the vendored NeoGPU fused 4-plane lm_head kernel (#1150). + * + * Wraps the bundled C symbol + * + * void skainet_ternary_lmhead_stage1( + * const float* input, int32_t input_offset, + * const uint8_t* planes, int32_t planes_byte_offset, int32_t plane_stride_bytes, + * const uint16_t* row_scale, int32_t row_scale_offset, + * int32_t input_dim, int32_t output_dim, + * float* output, int32_t output_offset); + * + * One call computes four planes with fused weights {1, ⅓, ⅑, 1/27} and applies the FP16 row + * scales; the pack's view kernel makes two calls for the full 8-plane result. The row-scale + * pointer is the weight segment sliced at [TernaryLmheadNative.lmheadStage1]'s + * `rowScaleByteOffset` (must be 2-byte aligned — it always is, plane strides are `rows·cols/4` + * with `cols % 4 == 0`). The C side threads internally with pthreads at any output_dim. + * + * Copy-in/copy-out per call like every FFM kernel here — fine for benches and small heads; + * a persistent off-heap weight arena is the follow-up for a 128k-vocab lm_head. + */ +@OptIn(ExperimentalMemoryApi::class) +public object NativeTernaryLmheadKernel : TernaryLmheadNative { + + override val name: String get() = "ffm" + + public fun isAvailable(): Boolean = handle != null + + /** Register with [TernaryPlanesKernelPack] when the bundled library resolves. */ + public fun install(warn: (String) -> Unit = {}): String = + TernaryPlanesKernelPack.install(if (isAvailable()) this else null, warn = warn) + + override fun lmheadStage1( + activation: FloatArray, activationOffset: Int, + weight: ByteArray, planesByteOffset: Int, + planeStrideBytes: Int, rowScaleByteOffset: Int, + inputDim: Int, outputDim: Int, + out: FloatArray, outOffset: Int, + ) { + require(inputDim % 4 == 0) { + "NativeTernaryLmheadKernel: inputDim must be a multiple of 4; got $inputDim" + } + require(rowScaleByteOffset % 2 == 0) { + "NativeTernaryLmheadKernel: rowScaleByteOffset must be 2-byte aligned; got $rowScaleByteOffset" + } + if (outputDim == 0) return + + val mh = handle + ?: error("NativeTernaryLmheadKernel invoked while native library unavailable") + + val inputReachFloats = if (inputDim == 0) 0 else activationOffset + inputDim + val outputReachFloats = outOffset + outputDim + + Arena.ofConfined().use { arena -> + val fAlign = ValueLayout.JAVA_FLOAT.byteAlignment() + val inputSeg: MemorySegment = if (inputReachFloats > 0) + arena.allocate(inputReachFloats.toLong() * java.lang.Float.BYTES, fAlign) + else MemorySegment.NULL + val weightSeg: MemorySegment = arena.allocate(weight.size.toLong(), ValueLayout.JAVA_SHORT.byteAlignment()) + val outputSeg: MemorySegment = + arena.allocate(outputReachFloats.toLong() * java.lang.Float.BYTES, fAlign) + + if (inputReachFloats > 0) { + MemorySegment.copy(activation, 0, inputSeg, ValueLayout.JAVA_FLOAT, 0L, inputReachFloats) + } + MemorySegment.copy(weight, 0, weightSeg, ValueLayout.JAVA_BYTE, 0L, weight.size) + // Round-trip the output reach so content before outOffset survives copy-back. + MemorySegment.copy(out, 0, outputSeg, ValueLayout.JAVA_FLOAT, 0L, outputReachFloats) + + mh.invoke( + inputSeg, activationOffset, + weightSeg, planesByteOffset, planeStrideBytes, + weightSeg.asSlice(rowScaleByteOffset.toLong()), 0, + inputDim, outputDim, + outputSeg, outOffset, + ) + + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, out, 0, outputReachFloats) + } + } + + private val handle: MethodHandle? by lazy { + val lookup = NativeLibraryLoader.lookup() ?: return@lazy null + val symbol = lookup.find("skainet_ternary_lmhead_stage1").orElse(null) ?: return@lazy null + val descriptor = FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, // input + ValueLayout.JAVA_INT, // input_offset + ValueLayout.ADDRESS, // planes + ValueLayout.JAVA_INT, // planes_byte_offset + ValueLayout.JAVA_INT, // plane_stride_bytes + ValueLayout.ADDRESS, // row_scale + ValueLayout.JAVA_INT, // row_scale_offset + ValueLayout.JAVA_INT, // input_dim + ValueLayout.JAVA_INT, // output_dim + ValueLayout.ADDRESS, // output + ValueLayout.JAVA_INT, // output_offset + ) + runCatching { Linker.nativeLinker().downcallHandle(symbol, descriptor) }.getOrNull() + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryPlanesFfmTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryPlanesFfmTest.kt new file mode 100644 index 000000000..34d5031fc --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/TernaryPlanesFfmTest.kt @@ -0,0 +1,67 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.KernelDispatch +import sk.ainet.backend.api.kernel.TernaryPlanesMatmulKernel +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.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.BitNetPlanesTensorData +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1150 end-to-end on the JVM: the REAL vendored `hs_ml_lmhead_stage1` (via + * `skainet_ternary_lmhead_stage1`, FFM) behind the REAL dispatcher, against the full 8-plane + * Kotlin reference. The C kernel spawns its 4 pthreads at any output_dim, so every case here + * also exercises its internal threading. + */ +@OptIn(ExperimentalMemoryApi::class) +class TernaryPlanesFfmTest { + + @BeforeTest fun setUp() { + KernelDispatch.clearForTesting() + assertTrue(NativeTernaryLmheadKernel.isAvailable(), "bundled libskainet_kernels must resolve") + } + @AfterTest fun tearDown() = KernelDispatch.clearForTesting() + + private fun assertDispatchParity(n: Int, k: Int, seed: Int) { + assertEquals("ternary_planes_matmul/ffm", NativeTernaryLmheadKernel.install()) + + val rng = Random(seed) + val values = FloatArray(n * k) { (rng.nextFloat() - 0.5f) * 2f } + val w = BitNetPlanesTensorData.fromFloats(Shape(n, k), values).packedView + val a = TensorView.dense( + Storage.Heap.wrap(FloatArray(k) { rng.nextFloat() - 0.5f }), + 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_planes_matmul/ffm", sink.eventsOf().single().kernel) + + val fromReference = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + TernaryPlanesMatmulKernel(TernaryPlanesMatmulKernel.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 (n=$n k=$k)", + ) + } + } + + @Test fun small_head() = assertDispatchParity(n = 8, k = 64, seed = 1) + + @Test fun bitnet_hidden_size_vocab_slice() = assertDispatchParity(n = 512, k = 2560, seed = 2) +} diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index c66fe837a..9c7b5dfb9 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -145,6 +145,45 @@ public class StreamingGgufParametersLoader( /** The dense FP32 size of [elements] — what every widening in this loader converts to. */ private fun denseFp32Bytes(elements: Long): Long = elements * 4 + /** Whether [form] asks for the one requantization this loader has (#1150). */ + private fun planesRequested(form: WeightForm): Boolean = + (form.encoding as? EncodingRequest.RequantizeTo)?.encoding == + sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_PLANES + + /** + * Requantize [values] (row-major `[out, in]` — validateForm pinned the orientation) into a + * packed [sk.ainet.lang.tensor.data.BitNetPlanesTensorData] (#1150): 8 trit planes + FP16 + * per-row scales, the lm_head format the planes kernel pack serves. Traced — a requantize is + * a conversion someone should be able to see. + */ + @Suppress("UNCHECKED_CAST") + private fun planesTensor( + ctx: ExecutionContext, + dtype: KClass, + shape: Shape, + tensorName: String, + values: FloatArray, + bytesBefore: Long, + fromFormat: Format, + ): Tensor { + require(shape.rank == 2) { + "tensor '$tensorName': RequantizeTo(BITNET_PLANES) needs a 2-D weight, got rank ${shape.rank}" + } + require(dtype == FP32::class) { + "tensor '$tensorName': BITNET_PLANES weights are logically FP32; requested $dtype" + } + val data = sk.ainet.lang.tensor.data.BitNetPlanesTensorData.fromFloats(shape, values) + traceConversion( + kind = "requantize-planes", + tensorName = tensorName, + from = fromFormat, + to = Format(FP32, sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_PLANES), + bytesBefore = bytesBefore, + bytesAfter = data.packedData.size.toLong(), + ) + return ctx.fromData(data as sk.ainet.lang.tensor.data.TensorData, dtype) + } + /** The uniform form: [weightForm] if given, else the historical as-stored-on-heap default. */ private val form: WeightForm = weightForm ?: WeightForm.AS_STORED_ON_HEAP @@ -169,10 +208,18 @@ public class StreamingGgufParametersLoader( "has no answer while the tensor is still labelled in the file's `ne` order." } val requested = form.encoding - require(requested !is EncodingRequest.RequantizeTo) { - "$where: EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " + - "is not supported by this loader: re-quantizing a weight the file does not already " + - "carry needs a quantizer per target encoding, and none of them exist here yet." + if (requested is EncodingRequest.RequantizeTo) { + // #1150: the one requantizer this loader has. Everything else still fails eagerly. + require(requested.encoding == sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_PLANES) { + "$where: EncodingRequest.RequantizeTo(${requested.encoding.name}) is not supported " + + "by this loader: re-quantizing a weight the file does not already carry needs a " + + "quantizer per target encoding, and only BITNET_PLANES exists here (#1150)." + } + require(form.shape == WeightShapeOrientation.OUT_IN) { + "$where: RequantizeTo(BITNET_PLANES) needs WeightShapeOrientation.OUT_IN — the " + + "format's scales are per output row, which has no answer while the tensor is " + + "still labelled in the file's `ne` order." + } } val dequantTarget = (requested as? EncodingRequest.DequantizeTo)?.dtype require(dequantTarget == null || dequantTarget == FP32) { @@ -252,7 +299,14 @@ public class StreamingGgufParametersLoader( when (dtype) { // The freshly decoded array is loader-owned — wrap it zero-copy // instead of paying the factory's defensive copy (#782). - FP32::class -> ctx.wrapFloatArray(shape, dtype, bytesToFloatArray(rawBytes)) as Tensor + FP32::class -> if (planesRequested(tensorForm)) { + planesTensor( + ctx, dtype, shape, tensorInfo.name, bytesToFloatArray(rawBytes), + rawBytes.size.toLong(), Format.dense(FP32), + ) + } else { + ctx.wrapFloatArray(shape, dtype, bytesToFloatArray(rawBytes)) as Tensor + } else -> null } } @@ -274,11 +328,18 @@ public class StreamingGgufParametersLoader( ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) } else { // Loader-owned widened array — zero-copy wrap (#782). - traceConversion( - "widen-f16", tensorInfo.name, Format.dense(FP16), Format.dense(FP32), - rawBytes.size.toLong(), denseFp32Bytes(tensorInfo.nElements), - ) - ctx.wrapFloatArray(shape, dtype, dequantF16(rawBytes)) as Tensor + if (planesRequested(tensorForm)) { + planesTensor( + ctx, dtype, shape, tensorInfo.name, dequantF16(rawBytes), + rawBytes.size.toLong(), Format.dense(FP16), + ) + } else { + traceConversion( + "widen-f16", tensorInfo.name, Format.dense(FP16), Format.dense(FP32), + rawBytes.size.toLong(), denseFp32Bytes(tensorInfo.nElements), + ) + ctx.wrapFloatArray(shape, dtype, dequantF16(rawBytes)) as Tensor + } } else -> null } @@ -356,6 +417,13 @@ public class StreamingGgufParametersLoader( rawBytes: ByteArray, tensorForm: WeightForm, ): Tensor { + if (planesRequested(tensorForm) && dtype == FP32::class) { + val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) + return planesTensor( + ctx, dtype, shape, tensorInfo.name, dest, + rawBytes.size.toLong(), ggufFormat(tensorInfo.tensorType, rawBytes.size.toLong()), + ) + } if (tensorForm.encoding is EncodingRequest.DequantizeTo && (dtype == FP32::class || dtype == FP16::class) ) { @@ -501,6 +569,14 @@ public class StreamingGgufParametersLoader( bytesBefore = rawBytes.size.toLong(), bytesAfter = packedBytes.size.toLong(), ) + if (planesRequested(tensorForm)) { + return planesTensor( + ctx, dtype, shape, tensorInfo.name, + sk.ainet.lang.memory.TernaryCodec.decodeBitNet(packedBytes, tensorInfo.nElements.toInt()), + rawBytes.size.toLong(), + Format(FP32, sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_B1_58), + ) + } if (tensorForm.encoding is EncodingRequest.DequantizeTo) { val dest = sk.ainet.lang.memory.TernaryCodec.decodeBitNet(packedBytes, tensorInfo.nElements.toInt()) traceConversion( diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PlanesRequantizeLoadTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PlanesRequantizeLoadTest.kt new file mode 100644 index 000000000..a01010e23 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PlanesRequantizeLoadTest.kt @@ -0,0 +1,106 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.BitNetPlanesTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** + * #1150: `RequantizeTo(BITNET_PLANES)` — the loader's one requantizer. A trained FP32 lm_head + * weight (or an I2_S one) arrives as packed [BitNetPlanesTensorData], encoded once at load, + * reconstructing within the format's truncation bound. + */ +@OptIn(ExperimentalMemoryApi::class) +class PlanesRequantizeLoadTest { + + private val planesForm = WeightForm( + encoding = EncodingRequest.RequantizeTo(TensorEncoding.BITNET_PLANES), + shape = WeightShapeOrientation.OUT_IN, + ) + + private fun f32Tensor(name: String, out: Int, inDim: Int, values: FloatArray): SyntheticGguf.TestTensor { + val buf = ByteBuffer.allocate(values.size * 4).order(ByteOrder.LITTLE_ENDIAN) + values.forEach { buf.putFloat(it) } + // GGUF `ne` order is fastest-varying first: a [out, in] row-major weight declares [in, out]. + return SyntheticGguf.TestTensor( + name, GGMLQuantizationType.F32, values.size.toLong(), buf.array(), + dims = listOf(inDim.toLong(), out.toLong()), + ) + } + + private fun load(file: File, form: WeightForm): Map> { + val ctx = DefaultDataExecutionContext() + val loaded = mutableMapOf>() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file) }, + weightForm = form, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + @Test + fun f32WeightRequantizesToPackedPlanesWithinTheTruncationBound() { + val out = 6; val inDim = 32 + val rng = Random(3) + val values = FloatArray(out * inDim) { (rng.nextFloat() - 0.5f) * 2f } + val file = SyntheticGguf.write(f32Tensor("output.weight", out, inDim, values)) + try { + val data = assertIs(load(file, planesForm).getValue("output.weight").data) + assertEquals(out, data.rows) + assertEquals(inDim, data.cols) + val decoded = TernaryCodec.decodeBitNetPlanes(data.packedData, out, inDim) + for (r in 0 until out) { + val bound = data.rowScale(r) * (0.5f / 2187f) + 1e-4f + for (c in 0 until inDim) { + val err = abs(values[r * inDim + c] - decoded[r * inDim + c]) + assertTrue(err <= bound, "[$r,$c]: err=$err > $bound") + } + } + } finally { + file.delete() + } + } + + @Test + fun requantizeToAnythingElseStillFailsEagerly() { + assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { throw IllegalStateException("never opened") }, + weightForm = WeightForm( + encoding = EncodingRequest.RequantizeTo(TensorEncoding.TQ2_0), + shape = WeightShapeOrientation.OUT_IN, + ), + ) + } + } + + @Test + fun planesWithoutOutInOrientationFailsEagerly() { + assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { throw IllegalStateException("never opened") }, + weightForm = WeightForm(encoding = EncodingRequest.RequantizeTo(TensorEncoding.BITNET_PLANES)), + ) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/BlockSpec.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/BlockSpec.kt index bafbd39f9..fe6ab2ca9 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/BlockSpec.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/BlockSpec.kt @@ -94,6 +94,12 @@ public val TensorEncoding.blockSpec: BlockSpec? BlockSpec(BlockSpec.PER_TENSOR_BLOCK, 0, 2.0, ScalePlacement.NONE, BlockSpec.INT8_ACTIVATION) TensorEncoding.BITNET_B1_58 -> BlockSpec(BlockSpec.PER_TENSOR_BLOCK, 0, 2.0, ScalePlacement.PER_TENSOR, BlockSpec.INT8_ACTIVATION) + // Deliberately no BlockSpec and — unlike the other ternary encodings — NO int8 activation + // hint: BITNET_PLANES is the f32-activation lm_head format (#1150); its geometry is + // row-scoped ([rows] FP16 scales after 8 plane payloads), which the row-count-free + // BlockSpec model cannot express, and without the planes kernel pack the dispatcher must + // fall to the decoding reference matmul, never the int8 requantize adapter. + TensorEncoding.BITNET_PLANES -> null is TensorEncoding.TurboQuantPolar -> BlockSpec(blockSize, (physicalBytes(blockSize.toLong()) ?: 0L).toInt(), bitsPerElement.toDouble(), ScalePlacement.BLOCK_HEAD) is TensorEncoding.TurboQuantPolarQjl -> diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TernaryCodec.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TernaryCodec.kt index f4d43c88a..669eb9f15 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TernaryCodec.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TernaryCodec.kt @@ -185,6 +185,90 @@ public object TernaryCodec { public fun bitNetScale(bytes: ByteArray, elementCount: Int, byteOffset: Int = 0): Float = Float.fromBits(le32(bytes, byteOffset + (elementCount + 3) / 4)) + // --- BitNet multi-plane trit residual (BITNET_PLANES, #1150) ------------------------------ + + /** + * Encode a `[rows, cols]` weight as [TensorEncoding.BITNET_PLANES] — the Kotlin port of + * NeoGPU's `hs_mlt_lmhead_encode`: per row, `scale = max|row|` stored as FP16; the row is + * normalized and decomposed by repeated round-to-trit (±0.5 threshold) with ×3 residual + * scaling into [TensorEncoding.BITNET_PLANES.PLANES] planes, each sequentially 2-bit-packed. + */ + public fun encodeBitNetPlanes(values: FloatArray, rows: Int, cols: Int): ByteArray { + require(cols % 4 == 0) { "BITNET_PLANES needs cols % 4 == 0; got $cols" } + require(values.size == rows * cols) { "values (${values.size}) must be rows*cols (${rows * cols})" } + val planes = TensorEncoding.BITNET_PLANES.PLANES + val planeStride = TensorEncoding.BITNET_PLANES.planeStrideBytes(rows, cols) + val scalesOffset = TensorEncoding.BITNET_PLANES.rowScalesByteOffset(rows, cols) + val out = ByteArray(TensorEncoding.BITNET_PLANES.bufferBytes(rows, cols)) + val rowBytes = cols / 4 + val residual = FloatArray(cols) + for (r in 0 until rows) { + var scale = 0f + for (c in 0 until cols) { + val a = abs(values[r * cols + c]) + if (a > scale) scale = a + } + val scaleBits = sk.ainet.lang.types.Fp16Codec.encode(scale) + out[scalesOffset + r * 2] = (scaleBits and 0xFF).toByte() + out[scalesOffset + r * 2 + 1] = ((scaleBits shr 8) and 0xFF).toByte() + // Decompose against the FP16-rounded scale the decoder will read, not the exact one — + // otherwise the stored trits answer for a scale that no longer exists. + val storedScale = sk.ainet.lang.types.Fp16Codec.decode(scaleBits) + val inv = if (storedScale != 0f) 1f / storedScale else 0f + for (c in 0 until cols) residual[c] = values[r * cols + c] * inv + for (p in 0 until planes) { + val base = p * planeStride + r * rowBytes + for (c in 0 until cols) { + val v = residual[c] + val t = if (v > 0.5f) 1 else if (v < -0.5f) -1 else 0 + val shift = (c % 4) * 2 + out[base + c / 4] = (out[base + c / 4].toInt() or ((t + 1) shl shift)).toByte() + residual[c] = (v - t) * 3f + } + } + } + return out + } + + /** Decode a full [TensorEncoding.BITNET_PLANES] buffer back to `rows*cols` floats (all planes). */ + public fun decodeBitNetPlanes(bytes: ByteArray, rows: Int, cols: Int, byteOffset: Int = 0): FloatArray { + val out = FloatArray(rows * cols) + for (r in 0 until rows) decodeBitNetPlanesRow(bytes, rows, cols, r, out, r * cols, byteOffset) + return out + } + + /** Decode one row of a [TensorEncoding.BITNET_PLANES] buffer into [out] at [outOffset]. */ + public fun decodeBitNetPlanesRow( + bytes: ByteArray, + rows: Int, + cols: Int, + row: Int, + out: FloatArray, + outOffset: Int, + byteOffset: Int = 0, + ) { + val planeStride = TensorEncoding.BITNET_PLANES.planeStrideBytes(rows, cols) + val rowBytes = cols / 4 + val scale = planesRowScale(bytes, rows, cols, row, byteOffset) + for (c in 0 until cols) { + var acc = 0f + var w = 1f + for (p in 0 until TensorEncoding.BITNET_PLANES.PLANES) { + val b = bytes[byteOffset + p * planeStride + row * rowBytes + c / 4].toInt() and 0xFF + acc += (((b shr ((c % 4) * 2)) and 3) - 1) * w + w /= 3f + } + out[outOffset + c] = acc * scale + } + } + + /** The FP16 per-row scale of row [row] in a `[rows, cols]` [TensorEncoding.BITNET_PLANES] buffer. */ + public fun planesRowScale(bytes: ByteArray, rows: Int, cols: Int, row: Int, byteOffset: Int = 0): Float { + val offset = byteOffset + TensorEncoding.BITNET_PLANES.rowScalesByteOffset(rows, cols) + row * 2 + val bits = (bytes[offset].toInt() and 0xFF) or ((bytes[offset + 1].toInt() and 0xFF) shl 8) + return sk.ainet.lang.types.Fp16Codec.decode(bits) + } + // --- dispatch ---------------------------------------------------------------------------- /** diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesTensorData.kt new file mode 100644 index 000000000..7fa1fdae1 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesTensorData.kt @@ -0,0 +1,82 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.DType + +/** + * Packed `[rows, cols]` weight in the [TensorEncoding.BITNET_PLANES] layout (#1150): 8 trit + * planes + FP16 per-row scales, decode pinned to [TernaryCodec.decodeBitNetPlanesRow]. + * + * One block per **row** ([blockSize] = cols): the format's scale granularity is the row, and the + * fused lm_head kernel consumes whole rows. [packedView] carries + * `Format(FP32, BITNET_PLANES) × BLOCKED_ROW_MAJOR` — the exact key the planes kernel pack + * serves. `get` returns the fully decoded value (all 8 planes × row scale); there is no single + * meaningful "code" per element in a residual format. + */ +public class BitNetPlanesTensorData( + initialShape: Shape, + private val data: ByteArray, +) : TensorData, PackedBlockStorage { + + /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ + @sk.ainet.lang.memory.ExperimentalMemoryApi + override val view: sk.ainet.lang.memory.TensorView get() = packedView + + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) + + public val rows: Int + public val cols: Int + + init { + require(shape.rank == 2) { "BITNET_PLANES is a [rows, cols] weight format; got rank ${shape.rank}" } + rows = shape[0] + cols = shape[1] + require(cols % 4 == 0) { "BITNET_PLANES needs cols % 4 == 0; got $cols" } + val required = TensorEncoding.BITNET_PLANES.bufferBytes(rows, cols) + require(data.size >= required) { + "BitNetPlanesTensorData: buffer is ${data.size} bytes, need >= $required for [$rows, $cols]" + } + } + + override val encoding: TensorEncoding get() = TensorEncoding.BITNET_PLANES + override val blockCount: Int get() = rows + override val blockSize: Int get() = cols + override val packedData: ByteArray get() = data + + /** The FP16 per-row scale of [row]. */ + public fun rowScale(row: Int): Float = TernaryCodec.planesRowScale(data, rows, cols, row) + + override fun dequantizeBlock(blockIdx: Int, output: FloatArray, outputOffset: Int) { + require(blockIdx in 0 until rows) { "row $blockIdx out of bounds (0..<$rows)" } + TernaryCodec.decodeBitNetPlanesRow(data, rows, cols, blockIdx, output, outputOffset) + } + + override fun get(vararg indices: Int): Float { + require(indices.size == 2) { "BITNET_PLANES data is 2-D" } + val row = FloatArray(cols) + TernaryCodec.decodeBitNetPlanesRow(data, rows, cols, indices[0], row, 0) + return row[indices[1]] + } + + override fun set(vararg indices: Int, value: Float) { + throw UnsupportedOperationException( + "BITNET_PLANES is a residual format — single elements cannot be re-encoded in place; " + + "re-encode the tensor with TernaryCodec.encodeBitNetPlanes", + ) + } + + public companion object { + /** Wrap raw plane bytes (validates the size against the shape). */ + public fun fromRawBytes(shape: Shape, bytes: ByteArray): BitNetPlanesTensorData = + BitNetPlanesTensorData(shape, bytes) + + /** Encode [values] (row-major `[rows, cols]`) with [TernaryCodec.encodeBitNetPlanes]. */ + public fun fromFloats(shape: Shape, values: FloatArray): BitNetPlanesTensorData { + require(shape.rank == 2) { "BITNET_PLANES is a [rows, cols] weight format" } + return BitNetPlanesTensorData(shape, TernaryCodec.encodeBitNetPlanes(values, shape[0], shape[1])) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorEncoding.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorEncoding.kt index 87147fff8..6b3bff89c 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorEncoding.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorEncoding.kt @@ -182,6 +182,46 @@ public sealed interface TensorEncoding { override fun physicalBytes(elementCount: Long): Long = (elementCount + 3) / 4 + SCALE_BYTES } + /** + * NeoGPU's multi-plane trit residual format for lm_head-class `[rows, cols]` weights (#1150): + * **8 sequentially-packed ternary planes + one FP16 scale per row**. + * + * Per row, `scale = max|row|` (stored as FP16); the normalized row is decomposed by repeated + * round-to-trit with ×3 residual scaling, so plane `p` carries weight `1/3^p` and + * + * `w[r, c] ≈ rowScale[r] · Σ_p (code_p(r, c) − 1) / 3^p` + * + * — effectively "16 bits as eight ternary digits", with truncation error ≤ `rowScale / (2·3⁷)`. + * Each plane is a full `[rows, cols]` matrix in the [BITNET_B1_58] payload order (4 codes per + * byte, low bit-pair first), `rows · cols/4` bytes, planes concatenated, then `rows` FP16 + * little-endian scales. `cols % 4 == 0` required. + * + * The point is *speed*, not memory (2 B + ε per weight — FP16-sized): the fused 4-plane LUT + * kernel reads planes 0–3 in one pass on baseline NEON, and an application can rescore top + * candidates with planes 4–7 (NeoGPU's two-stage lm_head). [physicalBytes] is `null` — the + * layout depends on the row count, not just the element count. + */ + public data object BITNET_PLANES : TensorEncoding { + /** Number of trit planes. */ + public const val PLANES: Int = 8 + + /** Bytes of FP16 per-row scale, per row, after the plane payloads. */ + public const val ROW_SCALE_BYTES: Int = 2 + + override val name: String get() = "BitNet-planes" + override fun physicalBytes(elementCount: Long): Long? = null + + /** Bytes of one plane of a `[rows, cols]` weight (`cols % 4 == 0`). */ + public fun planeStrideBytes(rows: Int, cols: Int): Int = rows * (cols / 4) + + /** Byte offset of the FP16 row-scale table inside the buffer. */ + public fun rowScalesByteOffset(rows: Int, cols: Int): Int = PLANES * planeStrideBytes(rows, cols) + + /** Total buffer size of a `[rows, cols]` weight. */ + public fun bufferBytes(rows: Int, cols: Int): Int = + rowScalesByteOffset(rows, cols) + rows * ROW_SCALE_BYTES + } + /** * TurboQuant PolarOnly encoding: rotation + scalar quantization + bit-packing. * diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesCodecTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesCodecTest.kt new file mode 100644 index 000000000..373c7014a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetPlanesCodecTest.kt @@ -0,0 +1,86 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.storage.TensorEncoding +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1150: the multi-plane trit residual codec (the Kotlin port of NeoGPU's `hs_mlt_lmhead_encode`). + * Eight planes of ±0.5-threshold round-to-trit with ×3 residual scaling reconstruct the row + * within `rowScale / (2·3⁷)` — the format's defined truncation bound. + */ +@OptIn(ExperimentalMemoryApi::class) +class BitNetPlanesCodecTest { + + @Test + fun eightPlanesReconstructWithinTheTruncationBound() { + val rows = 5; val cols = 64 + val rng = Random(7) + val values = FloatArray(rows * cols) { (rng.nextFloat() - 0.5f) * 2f } + val bytes = TernaryCodec.encodeBitNetPlanes(values, rows, cols) + val decoded = TernaryCodec.decodeBitNetPlanes(bytes, rows, cols) + for (r in 0 until rows) { + val scale = TernaryCodec.planesRowScale(bytes, rows, cols, r) + // bound: residual after 8 trits is at most 0.5/3^7 of the normalized value, plus + // FP16 rounding of the scale itself — allow a small epsilon on top. + val bound = scale * (0.5f / 2187f) + 1e-4f + for (c in 0 until cols) { + val err = abs(values[r * cols + c] - decoded[r * cols + c]) + assertTrue(err <= bound, "[$r,$c]: |${values[r * cols + c]} - ${decoded[r * cols + c]}| = $err > $bound") + } + } + } + + @Test + fun planeZeroCarriesTheSignStructure() { + val cols = 16 + val values = FloatArray(cols) { if (it % 3 == 0) 0.9f else if (it % 3 == 1) -0.9f else 0.0f } + val bytes = TernaryCodec.encodeBitNetPlanes(values, 1, cols) + for (c in 0 until cols) { + val code = ((bytes[c / 4].toInt() and 0xFF) shr ((c % 4) * 2)) and 3 + val expected = when (c % 3) { 0 -> 2; 1 -> 0; else -> 1 } // +1, -1, 0 biased + assertEquals(expected, code, "plane-0 trit of element $c") + } + } + + @Test + fun rowScaleIsTheRowsAbsMaxAsFp16() { + val values = floatArrayOf(0.1f, -0.75f, 0.5f, 0.25f, 0f, 0f, 0f, 0f) + val bytes = TernaryCodec.encodeBitNetPlanes(values, 1, 8) + assertEquals(0.75f, TernaryCodec.planesRowScale(bytes, 1, 8, 0), "0.75 is FP16-exact") + } + + @Test + fun tensorDataDecodesThroughTheSameCodec() { + val rows = 3; val cols = 32 + val rng = Random(11) + val values = FloatArray(rows * cols) { rng.nextFloat() - 0.5f } + val data = BitNetPlanesTensorData.fromFloats(Shape(rows, cols), values) + assertEquals(rows, data.blockCount) + assertEquals(cols, data.blockSize) + assertEquals(TensorEncoding.BITNET_PLANES, data.packedView.format.encoding) + val viaCodec = TernaryCodec.decodeBitNetPlanes(data.packedData, rows, cols) + for (r in 0 until rows) for (c in 0 until cols) { + assertEquals(viaCodec[r * cols + c], data.get(r, c), "[$r,$c]") + } + } + + @Test + fun bufferGeometryHelpersAgree() { + val n = 6; val k = 32 + val enc = TensorEncoding.BITNET_PLANES + assertEquals(n * k / 4, enc.planeStrideBytes(n, k)) + assertEquals(8 * n * k / 4, enc.rowScalesByteOffset(n, k)) + assertEquals(8 * n * k / 4 + 2 * n, enc.bufferBytes(n, k)) + assertEquals( + enc.bufferBytes(n, k), + TernaryCodec.encodeBitNetPlanes(FloatArray(n * k), n, k).size, + ) + } +}