From d1aea653b816e60b8650362431172e426371a8eb Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 14:41:03 +0200 Subject: [PATCH] feat(memory): int8 activation requant adapter and the reference bitnet_gemv MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1040 (SKEEP-003 P6, S2.8, proposal §5.3; M2-F3). A ternary weight is only cheap if the other operand is cheap too. #1033 gave the ternary encodings an `activation` hint; nothing produced that activation and nothing consumed it. - `TensorEncoding.DENSE_I8_ABSMAX`: int8 codes with a per-token absmax scale — `rows * cols` codes then `rows` FP32 scales. Deliberately not parameterized by the row length, because kernel selection keys on the Format and a per-shape encoding would make every hidden size its own key; the row count comes from the view's shape. `BlockSpec.INT8_ACTIVATION` now names this instead of the `Dense(1)` placeholder. - `I8Absmax.requantize(view, scope, sink)`: the adapter. Per-row absmax, scale = absmax/127, a zero row keeps scale 0 instead of dividing by it. It allocates in the caller's scope and emits `AdapterInserted` — the cost is real every step, so it is visible rather than hidden in a kernel. The view decodes through `get()` like every other packed format (rule 4). - `BitNetGemvKernel`: the reference `bitnet_gemv` — no multiplies in the inner loop, since a ternary weight is an add, a subtract, or nothing. Scales are factored out per block exactly as a NEON kernel (#1041) will factor them, and the weight's codes are decoded once per call rather than once per row. Registered for TQ1_0, TQ2_0 and BitNet b1.58. - `KernelDispatch` reads the weight encoding's `activation` hint: when a kernel exists for the requantized pair it inserts the adapter and selects it. Only ternary formats ask, so nothing else changes path. - `TernaryBlockDecoder(encoding, elementCount)` completes #1033's per-tensor case: BitNet b1.58's single scale covers the whole tensor, so the tensor is one block. Acceptance: the kernel matches the decoded definition for all three encodings within 1e-3 relative; an all-zero weight yields exactly zero; requantization stays within one scale step per element; the dispatcher selects `bitnet_gemv/reference` for a ternary weight and emits exactly one visible `requantize-i8-absmax` adapter, while a dense weight gets none. The adapter's price is asserted: `rows * cols + rows * 4` bytes — 4100 for one token of a 4096-wide model, the ≈ 4 KB/step §5.3 predicts. One test assertion is relative rather than exact because Kotlin/JS computes Float arithmetic in double precision, so `amax / 127` can differ in the last bit between the test and the implementation. Gate: scripts/pr-gate.sh — all legs passed (JVM, JS, Wasm, linuxX64). Co-Authored-By: Claude Fable 5 --- .../backend/api/kernel/BitNetGemvKernel.kt | 123 ++++++++++ .../backend/api/kernel/KernelDispatch.kt | 27 +++ .../backend/api/kernel/BitNetGemvTest.kt | 213 ++++++++++++++++++ .../api/jvm/skainet-lang-core.api | 25 ++ .../kotlin/sk/ainet/lang/memory/BlockSpec.kt | 11 +- .../kotlin/sk/ainet/lang/memory/I8Absmax.kt | 157 +++++++++++++ .../sk/ainet/lang/memory/TernaryCodec.kt | 52 ++++- .../lang/tensor/storage/TensorEncoding.kt | 21 ++ .../ainet/lang/memory/TernaryEncodingTest.kt | 2 +- 9 files changed, 617 insertions(+), 14 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/BitNetGemvKernel.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/I8Absmax.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/BitNetGemvKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/BitNetGemvKernel.kt new file mode 100644 index 000000000..f89f6af52 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/BitNetGemvKernel.kt @@ -0,0 +1,123 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.I8Absmax +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView +import sk.ainet.lang.memory.TernaryCodec +import sk.ainet.lang.memory.isTernary +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * The reference `bitnet_gemv` (SKEEP-003 §5.3, M2-F3): int8 activations against ternary weights, + * with **no multiplies in the inner loop**. + * + * A ternary weight is `-1`, `0` or `+1` times a scale, so the dot product is an add, a subtract or + * nothing at all — the whole point of 1.58-bit weights, and the shape a NEON kernel (#1041) will + * take. This one is plain Kotlin and runs everywhere: it is the thing the SIMD versions are + * checked against, so it is written for clarity, and the scales are factored out of the loop + * exactly as the vector kernels will factor them. + * + * Operands: `[rows, k]` activations in [I8Absmax.FORMAT] × `[n, k]` ternary weights in canonical + * (row-major block) order — the order [TernaryCodec] produces and GGUF stores. Output `[rows, n]`. + * + * The weight's codes are read once per call, not once per row: a decode step is one row against + * the whole matrix, so hoisting it is the difference between O(rows·n·k) decodes and O(n·k). + */ +@ExperimentalMemoryApi +public class BitNetGemvKernel(override val key: KernelKey) : ViewKernel { + + override val name: String get() = "bitnet_gemv/reference" + + override fun run(inputs: List, out: TensorView) { + require(inputs.size == 2) { "bitnet_gemv takes (activation, weight), got ${inputs.size} operands" } + val a = inputs[0] + val w = inputs[1] + require(a.format == I8Absmax.FORMAT) { "activation must be ${I8Absmax.FORMAT}, was ${a.format}" } + require(w.format.encoding.isTernary) { "weight must be ternary, was ${w.format}" } + require(a.shape.rank == 2 && w.shape.rank == 2 && out.shape.rank == 2) { "bitnet_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}" } + + val encoding = w.format.encoding + val bytes = weightBytes(w) + val codes = TernaryCodec.codes(encoding, bytes, n * k) + val blockSize = blockSizeOf(encoding, n * k) + + for (r in 0 until rows) { + val activation = I8Absmax.rowCodes(a, r) + val activationScale = I8Absmax.scaleOf(a, r) + for (o in 0 until n) { + var acc = 0f + var index = o * k + var offset = 0 + // Walk the row block by block: within a block the weight scale is constant, so the + // inner loop only adds and subtracts activation codes (and skips the zeros). + while (offset < k) { + val run = minOf(blockSize - (index % blockSize), k - offset) + var partial = 0 + for (i in 0 until run) { + when (codes[index + i].toInt()) { + 1 -> partial += activation[offset + i].toInt() + -1 -> partial -= activation[offset + i].toInt() + else -> Unit // zero weights cost nothing + } + } + acc += partial * blockScale(encoding, bytes, (index + run - 1) / blockSize) + index += run + offset += run + } + out.set(r, o, value = acc * activationScale) + } + } + } + + private fun weightBytes(w: TensorView): ByteArray { + val heap = w.storage as? Storage.Heap + ?: throw UnsupportedOperationException("bitnet_gemv reads ternary weights from heap storage in this milestone") + return heap.bytes ?: throw UnsupportedOperationException("ternary weights need byte storage") + } + + /** Elements per scale: a GGML block, or the whole tensor for the per-tensor BitNet encoding. */ + private fun blockSizeOf(encoding: TensorEncoding, elements: Int): Int = + when (encoding) { + TensorEncoding.TQ1_0 -> TensorEncoding.TQ1_0.BLOCK_SIZE + TensorEncoding.TQ2_0 -> TensorEncoding.TQ2_0.BLOCK_SIZE + else -> elements + } + + /** The scale of block [block] — per-block FP16 for the GGML types, one FP32 for BitNet. */ + private fun blockScale(encoding: TensorEncoding, bytes: ByteArray, block: Int): Float = when (encoding) { + TensorEncoding.TQ1_0 -> fp16At(bytes, block * TensorEncoding.TQ1_0.BYTES_PER_BLOCK + 52) + TensorEncoding.TQ2_0 -> fp16At(bytes, block * TensorEncoding.TQ2_0.BYTES_PER_BLOCK + 64) + else -> TernaryCodec.bitNetScale(bytes, (bytes.size - TensorEncoding.BITNET_B1_58.SCALE_BYTES) * 4) + } + + private fun fp16At(bytes: ByteArray, offset: Int): Float = + sk.ainet.lang.types.Fp16Codec.decode( + (bytes[offset].toInt() and 0xFF) or ((bytes[offset + 1].toInt() and 0xFF) shl 8), + ) + + public companion object { + /** The key this kernel serves for a ternary [weightFormat]. */ + public fun keyFor(weightFormat: Format): KernelKey = KernelKey( + op = "matmul", + operands = listOf( + OperandKey.contiguous(I8Absmax.FORMAT), + OperandKey(weightFormat, LayoutClass.BLOCKED), + ), + ) + + /** Register the reference kernel for every ternary encoding that carries its own bytes. */ + public fun registerReference() { + for (encoding in listOf(TensorEncoding.TQ1_0, TensorEncoding.TQ2_0, TensorEncoding.BITNET_B1_58)) { + val format = Format(sk.ainet.lang.types.FP32, encoding) + KernelDispatch.register(BitNetGemvKernel(keyFor(format))) + } + } + } +} 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 index 03e9cc038..f7cd6cef6 100644 --- 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 @@ -4,6 +4,7 @@ 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.blockSpec import sk.ainet.lang.memory.trace.NoopTraceSink import sk.ainet.lang.memory.trace.TraceEvent import sk.ainet.lang.memory.trace.TraceSink @@ -76,6 +77,21 @@ public object KernelDispatch { runTraced(exact, listOf(a, b), out, sink) return } + // The weight's encoding may *ask* for a different activation format — a ternary weight wants + // int8 with a per-token scale (`W1.58A8`, §5.3). Honour the request when a kernel exists for + // the requantized pair: the adapter costs bytes in the caller's scope every step, so it is + // allocated there and emitted as an AdapterInserted rather than hidden inside the kernel. + val wanted = b.format.encoding.blockSpec?.activation + if (wanted != null && wanted != a.format) { + val requantized = requantizeFor(wanted, a, scope, sink) + if (requantized != null) { + val ternaryKernel = find(KernelKey.matmul(requantized, b)) + if (ternaryKernel != null) { + runTraced(ternaryKernel, listOf(requantized, 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") @@ -83,6 +99,17 @@ public object KernelDispatch { runTraced(reference, listOf(adaptedA, b), out, sink) } + /** + * Convert [activation] into the [wanted] activation format, or `null` when no adapter for it + * exists. Today the only one is the int8 absmax requantization the ternary kernels ask for. + */ + private fun requantizeFor(wanted: Format, activation: TensorView, scope: Scope, sink: TraceSink): TensorView? = + if (wanted == sk.ainet.lang.memory.I8Absmax.FORMAT && activation.shape.rank == 2) { + sk.ainet.lang.memory.I8Absmax.requantize(activation, scope, sink) + } else { + null + } + /** 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 diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTest.kt new file mode 100644 index 000000000..861e59851 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/BitNetGemvTest.kt @@ -0,0 +1,213 @@ +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.I8Absmax +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 + +/** + * #1040 (M2-F3): int8 activations against ternary weights, and the adapter that produces them. + */ +@OptIn(ExperimentalMemoryApi::class) +class BitNetGemvTest { + + private val k = 256 // one TQ block + private val n = 4 // output rows + + @BeforeTest fun setUp() { + KernelDispatch.clearForTesting() + BitNetGemvKernel.registerReference() + } + + @AfterTest fun tearDown() = KernelDispatch.clearForTesting() + + /** Deterministic ternary weight values, `[n, k]`, scaled so the block absmax is exact in FP16. */ + private fun weightValues(seed: Int = 3): FloatArray { + var s = seed + return FloatArray(n * k) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 3 - 1) * 0.5f + } + } + + private fun activationValues(rows: Int, seed: Int = 11): FloatArray { + var s = seed + return FloatArray(rows * k) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 2000 - 1000) / 1000f + } + } + + private fun weightView(encoding: TensorEncoding, values: FloatArray): TensorView { + val bytes = TernaryCodec.encode(encoding, values) + val decoder = if (encoding == TensorEncoding.BITNET_B1_58) { + TernaryBlockDecoder(encoding, values.size) // one scale for the whole tensor + } else { + TernaryBlockDecoder(encoding) + } + return TensorView.packed(Storage.Heap.wrap(bytes), Shape(n, k), encoding, decoder) + } + + /** `out[r, o] = Σ_k decoded_activation[r, k] * decoded_weight[o, k]` — the definition. */ + private fun reference(activation: TensorView, weight: TensorView, rows: Int): FloatArray { + val out = FloatArray(rows * n) + for (r in 0 until rows) { + for (o in 0 until n) { + var acc = 0f + for (i in 0 until k) acc += I8Absmax.valueAt(activation, r, i) * weight.get(o, i) + out[r * n + o] = acc + } + } + return out + } + + // --- the adapter --------------------------------------------------------------------------- + + @Test + fun requantizationKeepsTheValuesAndPricesItself() { + val rows = 2 + val values = activationValues(rows) + val dense = TensorView.dense(Storage.Heap.wrap(values), Shape(rows, k), FP32) + val sink = RecordingTraceSink() + val quantized = I8Absmax.requantize(dense, Scope.Ambient, sink) + + assertEquals(I8Absmax.FORMAT, quantized.format) + for (r in 0 until rows) { + var amax = 0f + for (c in 0 until k) amax = maxOf(amax, abs(values[r * k + c])) + // Kotlin/JS computes Float arithmetic in double precision, so the same division can + // differ in the last bit between the test and the implementation: compare relatively. + val expectedScale = amax / 127f + assertTrue( + abs(expectedScale - I8Absmax.scaleOf(quantized, r)) <= 1e-6f * expectedScale, + "row $r scale should be absmax / 127 = $expectedScale, was ${I8Absmax.scaleOf(quantized, r)}", + ) + val tolerance = I8Absmax.scaleOf(quantized, r) + for (c in 0 until k) { + assertTrue( + abs(values[r * k + c] - I8Absmax.valueAt(quantized, r, c)) <= tolerance, + "row $r col $c: ${values[r * k + c]} vs ${I8Absmax.valueAt(quantized, r, c)}", + ) + } + } + + val adapter = sink.eventsOf().single() + assertEquals("requantize-i8-absmax", adapter.kind) + assertEquals(I8Absmax.bytesFor(rows, k), adapter.bytes, "codes plus one scale per row") + assertEquals(rows.toLong() * k + rows * 4, adapter.bytes) + } + + @Test + fun aZeroRowSurvivesQuantization() { + val dense = TensorView.dense(Storage.Heap.wrap(FloatArray(k)), Shape(1, k), FP32) + val quantized = I8Absmax.requantize(dense, Scope.Ambient) + assertEquals(0f, I8Absmax.scaleOf(quantized, 0), "no division by zero") + for (c in 0 until k) assertEquals(0f, I8Absmax.valueAt(quantized, row = 0, col = c)) + } + + @Test + fun theAdapterCostPerStepIsTheOneTheDesignPredicts() { + // §5.3 quotes ≈ 4 KB per decode step for a 2 B-parameter model (hidden 4096, one token). + val bytes = I8Absmax.bytesFor(rows = 1, cols = 4096) + assertEquals(4096L + 4, bytes) + assertTrue(bytes < 5 * 1024, "one token's activations must stay in the kilobytes: $bytes") + } + + // --- the kernel ---------------------------------------------------------------------------- + + @Test + fun theReferenceKernelMatchesTheDefinitionForEveryTernaryEncoding() { + for (encoding in listOf(TensorEncoding.TQ1_0, TensorEncoding.TQ2_0, TensorEncoding.BITNET_B1_58)) { + val rows = 2 + val weight = weightView(encoding, weightValues()) + val activation = I8Absmax.requantize( + TensorView.dense(Storage.Heap.wrap(activationValues(rows)), Shape(rows, k), FP32), + Scope.Ambient, + ) + val out = TensorView.dense(Storage.Heap.floats(rows * n), Shape(rows, n), FP32) + BitNetGemvKernel(BitNetGemvKernel.keyFor(weight.format)).run(listOf(activation, weight), out) + + val expected = reference(activation, weight, rows) + for (r in 0 until rows) for (o in 0 until n) { + val got = out.get(r, o) + val want = expected[r * n + o] + assertTrue( + abs(got - want) <= 1e-3f * maxOf(1f, abs(want)), + "${encoding.name} [$r,$o]: $got vs $want", + ) + } + } + } + + @Test + fun zeroWeightsContributeNothing() { + // every weight zero → every output zero, whatever the activations are + val encoding = TensorEncoding.TQ2_0 + val weight = weightView(encoding, FloatArray(n * k)) + val activation = I8Absmax.requantize( + TensorView.dense(Storage.Heap.wrap(activationValues(1)), Shape(1, k), FP32), + Scope.Ambient, + ) + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + BitNetGemvKernel(BitNetGemvKernel.keyFor(weight.format)).run(listOf(activation, weight), out) + for (o in 0 until n) assertEquals(0f, out.get(0, o)) + } + + // --- dispatch ------------------------------------------------------------------------------ + + @Test + fun theDispatcherRequantizesTheActivationAndPicksTheTernaryKernel() { + val encoding = TensorEncoding.TQ2_0 + val weight = weightView(encoding, weightValues()) + val floats = activationValues(1) + val activation = TensorView.dense(Storage.Heap.wrap(floats), Shape(1, k), FP32) + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + + val sink = RecordingTraceSink() + val scope = ForwardScope(slabFloats = 4 * k, sink = sink, name = "decode") + KernelDispatch.matmul(activation, weight, out, scope, sink) + + val kernels = sink.eventsOf() + assertEquals(1, kernels.size) + assertEquals("bitnet_gemv/reference", kernels.single().kernel, "a ternary weight selects the ternary kernel") + val adapter = sink.eventsOf().single() + assertEquals("requantize-i8-absmax", adapter.kind, "and the activation adapter is visible, not hidden") + assertEquals(I8Absmax.FORMAT, adapter.to) + + // the numbers are the kernel's own + val quantized = I8Absmax.requantize(activation, Scope.Ambient) + val expected = reference(quantized, weight, rows = 1) + for (o in 0 until n) assertTrue(abs(out.get(0, o) - expected[o]) <= 1e-3f * maxOf(1f, abs(expected[o])), "[$o]") + scope.close() + } + + @Test + fun aDenseWeightIsUntouchedByAnyOfThis() { + val weightFloats = FloatArray(n * k) { (it % 5) * 0.25f } + val weight = TensorView.dense(Storage.Heap.wrap(weightFloats), Shape(n, k), FP32) + val activation = TensorView.dense(Storage.Heap.wrap(activationValues(1)), Shape(1, k), FP32) + val out = TensorView.dense(Storage.Heap.floats(n), Shape(1, n), FP32) + val sink = RecordingTraceSink() + KernelDispatch.matmul(activation, weight, out, Scope.Ambient, sink) + assertTrue( + sink.eventsOf().none { it.kind == "requantize-i8-absmax" }, + "only ternary formats ask for int8 activations", + ) + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index ec6247077..6248faeb6 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -795,6 +795,20 @@ public final class sk/ainet/lang/memory/ForwardScope : sk/ainet/lang/memory/Scop public static synthetic fun retain$default (Lsk/ainet/lang/memory/ForwardScope;Lsk/ainet/lang/memory/Storage$Heap;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/Storage$Heap; } +public final class sk/ainet/lang/memory/I8Absmax { + public static final field INSTANCE Lsk/ainet/lang/memory/I8Absmax; + public final fun bytesFor (II)J + public final fun codeAt (Lsk/ainet/lang/memory/TensorView;II)I + public final fun getFORMAT ()Lsk/ainet/lang/memory/Format; + public final fun requantize (Lsk/ainet/lang/memory/TensorView;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/memory/trace/TraceSink;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun requantize$default (Lsk/ainet/lang/memory/I8Absmax;Lsk/ainet/lang/memory/TensorView;Lsk/ainet/lang/memory/Scope;Lsk/ainet/lang/memory/trace/TraceSink;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; + public final fun rowCodes (Lsk/ainet/lang/memory/TensorView;I)[B + public final fun scaleOf (Lsk/ainet/lang/memory/TensorView;I)F + public final fun valueAt (Lsk/ainet/lang/memory/TensorView;II)F + public final fun view (Lsk/ainet/lang/memory/Storage;IILsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun view$default (Lsk/ainet/lang/memory/I8Absmax;Lsk/ainet/lang/memory/Storage;IILsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; +} + public final class sk/ainet/lang/memory/Layout { public static final field Companion Lsk/ainet/lang/memory/Layout$Companion; public fun (Lsk/ainet/lang/tensor/Shape;[IJIZI)V @@ -1209,6 +1223,7 @@ public final class sk/ainet/lang/memory/TensorView$Companion { public final class sk/ainet/lang/memory/TernaryBlockDecoder : sk/ainet/lang/memory/BlockDecoder { public fun (Lsk/ainet/lang/tensor/storage/TensorEncoding;)V + public fun (Lsk/ainet/lang/tensor/storage/TensorEncoding;I)V public fun decodeBlock (Lsk/ainet/lang/memory/Storage;J[FI)V public fun decodeElement (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/memory/Layout;J)F public fun getBlockSize ()I @@ -7131,6 +7146,16 @@ public final class sk/ainet/lang/tensor/storage/TensorEncoding$BITNET_B1_58 : sk public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/tensor/storage/TensorEncoding$DENSE_I8_ABSMAX : sk/ainet/lang/tensor/storage/TensorEncoding { + public static final field CODE_RANGE I + public static final field INSTANCE Lsk/ainet/lang/tensor/storage/TensorEncoding$DENSE_I8_ABSMAX; + public fun equals (Ljava/lang/Object;)Z + public fun getName ()Ljava/lang/String; + public fun hashCode ()I + public fun physicalBytes (J)Ljava/lang/Long; + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/tensor/storage/TensorEncoding$Dense : sk/ainet/lang/tensor/storage/TensorEncoding { public fun (I)V public final fun component1 ()I 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 792a42403..bafbd39f9 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 @@ -56,8 +56,12 @@ public data class BlockSpec( /** [blockSize] value meaning "the whole tensor is one block". */ public const val PER_TENSOR_BLOCK: Int = 0 - /** The activation format of the ternary kernels: int8, one byte per element (`W1.58A8`). */ - public val INT8_ACTIVATION: Format = Format(Int8, TensorEncoding.Dense(1)) + /** + * The activation format of the ternary kernels (`W1.58A8`): int8 codes with a per-token + * absmax scale — [TensorEncoding.DENSE_I8_ABSMAX], which is what the requant adapter of + * #1040 produces and what `bitnet_gemv` consumes. + */ + public val INT8_ACTIVATION: Format = Format(Int8, TensorEncoding.DENSE_I8_ABSMAX) } } @@ -74,6 +78,9 @@ public val TensorEncoding.blockSpec: BlockSpec? get() = when (this) { is TensorEncoding.Dense -> null is TensorEncoding.Opaque -> null + // Activations, not weights: the "block" is a row, whose length is the tensor's, not the + // encoding's — see the note on DENSE_I8_ABSMAX. + TensorEncoding.DENSE_I8_ABSMAX -> null TensorEncoding.Q4_0 -> BlockSpec(32, 18, 4.0, ScalePlacement.BLOCK_HEAD) TensorEncoding.Q5_0 -> BlockSpec(32, 22, 5.0, ScalePlacement.BLOCK_HEAD) TensorEncoding.Q5_1 -> BlockSpec(32, 24, 5.0, ScalePlacement.BLOCK_HEAD) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/I8Absmax.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/I8Absmax.kt new file mode 100644 index 000000000..2ad362f3c --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/I8Absmax.kt @@ -0,0 +1,157 @@ +package sk.ainet.lang.memory + +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.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.Int8 +import kotlin.math.abs +import kotlin.math.roundToInt + +/** + * The int8 activations the ternary kernels consume (`W1.58A8`, SKEEP-003 §5.3, M2-F3) and the + * adapter that produces them. + * + * A ternary weight is only cheap if the *other* operand is cheap too: `bitnet_gemv` adds and + * subtracts activation values by the sign of the weight, so the activations are quantized to int8 + * with one absmax scale per token. This is exactly the "dispatcher-inserted adapter" of §5.1 — it + * costs real bytes in the Forward scope every step, so it is allocated there and **traced**, never + * hidden inside a kernel. + * + * Byte layout of a `[rows, cols]` activation: `rows * cols` codes, row-major, then `rows` + * little-endian FP32 scales. `value ≈ code * scale(row)`. + */ +@ExperimentalMemoryApi +public object I8Absmax { + + /** The activation format: int8 values, per-token absmax scale. */ + public val FORMAT: Format = Format(Int8, TensorEncoding.DENSE_I8_ABSMAX) + + /** Bytes a `[rows, cols]` quantized activation occupies: the codes plus one scale per row. */ + public fun bytesFor(rows: Int, cols: Int): Long = rows.toLong() * cols + rows.toLong() * 4 + + /** + * Quantize [activation] (`[rows, cols]`, any decodable format) into [scope] as int8 codes with + * a per-row absmax scale, and emit the adapter event that prices it. + * + * A row of zeros keeps scale `0` and decodes back to zeros rather than dividing by zero. + */ + public fun requantize( + activation: TensorView, + scope: Scope, + sink: TraceSink = NoopTraceSink, + id: TensorId? = null, + ): TensorView { + require(activation.shape.rank == 2) { "activations are [rows, cols], got ${activation.shape}" } + val rows = activation.shape[0] + val cols = activation.shape[1] + val bytes = ByteArray(bytesFor(rows, cols).toInt()) + val scaleBase = rows * cols + for (r in 0 until rows) { + var amax = 0f + for (c in 0 until cols) amax = maxOf(amax, abs(activation.get(r, c))) + val scale = amax / TensorEncoding.DENSE_I8_ABSMAX.CODE_RANGE + val inverse = if (scale != 0f) 1f / scale else 0f + for (c in 0 until cols) { + val code = (activation.get(r, c) * inverse).roundToInt().coerceIn(-127, 127) + bytes[r * cols + c] = code.toByte() + } + putFloat(bytes, scaleBase + r * 4, scale) + } + val storage = Storage.Heap.wrap(bytes, mutable = false, origin = id ?: activation.id) + val target = id ?: activation.id + if (sink.isEnabled) { + sink.emit( + TraceEvent.AdapterInserted( + kind = "requantize-i8-absmax", + from = activation.format, + to = FORMAT, + bytes = bytesFor(rows, cols), + target = target, + scope = scope.kind, + ), + ) + } + return view(storage, rows, cols, target) + } + + /** A view over already-quantized bytes: `rows * cols` codes then `rows` FP32 scales. */ + public fun view(storage: Storage, rows: Int, cols: Int, id: TensorId? = null): TensorView { + val shape = Shape(rows, cols) + return TensorView( + shape = shape, + format = FORMAT, + layout = Layout(shape, Layout.rowMajorStrides(shape), 0L, 1), + storage = storage, + id = id, + decoder = Decoder(rows, cols), + ) + } + + /** + * Decodes an int8 code back to `code * scale(row)` — so `view.get(r, c)` returns a *value*, + * never a raw byte (rule 4), exactly like every other packed format. + */ + private class Decoder(private val rows: Int, private val cols: Int) : BlockDecoder { + override val blockSize: Int get() = 1 + override val bytesPerBlock: Int get() = 1 + + override fun decodeBlock(storage: Storage, blockIndex: Long, out: FloatArray, outOffset: Int) { + out[outOffset] = decode(storage, blockIndex) + } + + override fun decodeElement(storage: Storage, layout: Layout, flatElementIndex: Long): Float = + decode(storage, flatElementIndex) + + private fun decode(storage: Storage, flatIndex: Long): Float { + val heap = storage as? Storage.Heap + ?: throw UnsupportedOperationException("I8-absmax activations live on the heap in this milestone") + val bytes = heap.bytes ?: throw UnsupportedOperationException("I8-absmax activations need byte storage") + val index = flatIndex.toInt() + val row = index / cols + return bytes[heap.arrayOffset + index] * scaleAt(bytes, heap.arrayOffset + rows * cols + row * 4) + } + } + + private fun scaleAt(bytes: ByteArray, offset: Int): Float = Float.fromBits( + (bytes[offset].toInt() and 0xFF) or ((bytes[offset + 1].toInt() and 0xFF) shl 8) or + ((bytes[offset + 2].toInt() and 0xFF) shl 16) or ((bytes[offset + 3].toInt() and 0xFF) shl 24), + ) + + /** The int8 code at `(row, col)` of a view in [FORMAT]. */ + public fun codeAt(view: TensorView, row: Int, col: Int): Int { + val bytes = bytesOf(view) + return bytes[row * view.shape[1] + col].toInt() + } + + /** The absmax scale of [row]. */ + public fun scaleOf(view: TensorView, row: Int): Float = + scaleAt(bytesOf(view), view.shape[0] * view.shape[1] + row * 4) + + /** The decoded value at `(row, col)` — `code * scale(row)`. */ + public fun valueAt(view: TensorView, row: Int, col: Int): Float = codeAt(view, row, col) * scaleOf(view, row) + + /** The codes of [row] as a `ByteArray` view into the storage — what a kernel iterates. */ + public fun rowCodes(view: TensorView, row: Int): ByteArray { + val bytes = bytesOf(view) + val cols = view.shape[1] + return bytes.copyOfRange(row * cols, row * cols + cols) + } + + private fun bytesOf(view: TensorView): ByteArray { + require(view.format == FORMAT) { "not an I8-absmax activation: ${view.format}" } + val heap = view.storage as? Storage.Heap + ?: throw UnsupportedOperationException("I8-absmax activations live on the heap in this milestone") + return heap.bytes ?: throw UnsupportedOperationException("I8-absmax activations need byte storage") + } + + private fun putFloat(bytes: ByteArray, offset: Int, value: Float) { + val bits = value.toRawBits() + bytes[offset] = (bits and 0xFF).toByte() + bytes[offset + 1] = ((bits ushr 8) and 0xFF).toByte() + bytes[offset + 2] = ((bits ushr 16) and 0xFF).toByte() + bytes[offset + 3] = ((bits ushr 24) and 0xFF).toByte() + } +} 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 63b8bbb6e..f4d43c88a 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 @@ -307,18 +307,28 @@ public object TernaryCodec { * (SKEEP-003 §4.4: `get()` decodes, never a raw byte). */ @ExperimentalMemoryApi -public class TernaryBlockDecoder(private val encoding: TensorEncoding) : BlockDecoder { +public class TernaryBlockDecoder private constructor( + private val encoding: TensorEncoding, + override val blockSize: Int, + override val bytesPerBlock: Int, +) : BlockDecoder { + + /** A decoder for a block-structured ternary encoding ([TensorEncoding.TQ1_0], [TensorEncoding.TQ2_0]). */ + public constructor(encoding: TensorEncoding) : this( + encoding, + blockSize = requireBlocked(encoding).blockSize, + bytesPerBlock = requireBlocked(encoding).bytesPerBlock, + ) - private val spec: BlockSpec = encoding.blockSpec - ?: throw IllegalArgumentException("$encoding has no block spec") - - init { - require(encoding.isTernary) { "$encoding is not a ternary encoding" } - require(!spec.isPerTensor) { "${encoding.name} is a per-tensor encoding; decode it through its TensorData" } - } - - override val blockSize: Int get() = spec.blockSize - override val bytesPerBlock: Int get() = spec.bytesPerBlock + /** + * A decoder for a **per-tensor** ternary encoding ([TensorEncoding.BITNET_B1_58]), whose single + * scale covers all [elementCount] elements: the whole tensor is one block (#1040). + */ + public constructor(encoding: TensorEncoding, elementCount: Int) : this( + encoding, + blockSize = elementCount, + bytesPerBlock = requirePerTensor(encoding, elementCount), + ) override fun decodeBlock(storage: Storage, blockIndex: Long, out: FloatArray, outOffset: Int) { val heap = storage as? Storage.Heap @@ -327,4 +337,24 @@ public class TernaryBlockDecoder(private val encoding: TensorEncoding) : BlockDe val off = heap.arrayOffset + (blockIndex * bytesPerBlock).toInt() TernaryCodec.decode(encoding, bytes, blockSize, off).copyInto(out, outOffset) } + + private companion object { + fun requireBlocked(encoding: TensorEncoding): BlockSpec { + val spec = encoding.blockSpec ?: throw IllegalArgumentException("$encoding has no block spec") + require(encoding.isTernary) { "$encoding is not a ternary encoding" } + require(!spec.isPerTensor) { + "${encoding.name} is a per-tensor encoding; give this constructor the element count" + } + return spec + } + + fun requirePerTensor(encoding: TensorEncoding, elementCount: Int): Int { + val spec = encoding.blockSpec ?: throw IllegalArgumentException("$encoding has no block spec") + require(encoding.isTernary) { "$encoding is not a ternary encoding" } + require(spec.isPerTensor) { "${encoding.name} is block-structured; use the single-argument constructor" } + require(elementCount > 0) { "elementCount must be > 0" } + return (encoding.physicalBytes(elementCount.toLong()) + ?: throw IllegalArgumentException("${encoding.name} cannot size $elementCount elements")).toInt() + } + } } 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 5386148ef..87147fff8 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 @@ -263,6 +263,27 @@ public sealed interface TensorEncoding { } } + /** + * Int8 activations with a **per-row (per-token) absmax scale** — the companion format of the + * ternary weights (`W1.58A8`, #1040). + * + * Layout of a `[rows, cols]` activation: `rows * cols` int8 codes in row-major order, followed + * by `rows` little-endian FP32 scales. A value is `code * scale(row)`; the scale is + * `absmax(row) / 127`, so a row of zeros has scale zero and decodes to zeros. + * + * Deliberately not parameterized by the row length: kernel selection keys on the [Format], and + * a per-shape encoding would make every hidden size a different key. The row count comes from + * the view's shape, which is where it belongs. [physicalBytes] is therefore `null` — the byte + * count needs the row length, and [sk.ainet.lang.memory.I8Absmax.bytesFor] computes it. + */ + public data object DENSE_I8_ABSMAX : TensorEncoding { + /** Largest magnitude an int8 code may take; the scale is `absmax / this`. */ + public const val CODE_RANGE: Int = 127 + + override val name: String get() = "I8-absmax" + override fun physicalBytes(elementCount: Long): Long? = null + } + /** * Opaque / unknown encoding. Used as a fallback for formats the runtime * cannot yet interpret but still wants to carry through without error. diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TernaryEncodingTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TernaryEncodingTest.kt index 827c3b3d9..c47d6fd18 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TernaryEncodingTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/TernaryEncodingTest.kt @@ -55,7 +55,7 @@ class TernaryEncodingTest { fun theTernaryFamilyIsDescribedAsTernaryWithInt8Activations() { for (e in listOf(TensorEncoding.TQ1_0, TensorEncoding.TQ2_0, TensorEncoding.TernaryPacked, TensorEncoding.BITNET_B1_58)) { assertTrue(e.isTernary, "${e.name} must be ternary") - assertEquals(Format(Int8, TensorEncoding.Dense(1)), e.blockSpec?.activation, "${e.name}: W1.58A8") + assertEquals(Format(Int8, TensorEncoding.DENSE_I8_ABSMAX), e.blockSpec?.activation, "${e.name}: W1.58A8") } assertTrue(!TensorEncoding.Q4_K.isTernary) assertNull(TensorEncoding.Dense(4).blockSpec, "dense is not block-structured")