From b8e2f5c1bb8cf847a7309d234152e40a8886d160 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 13:38:04 +0200 Subject: [PATCH] =?UTF-8?q?feat(io):=20I2=5FS=20GGUF=20import=20=E2=80=94?= =?UTF-8?q?=20the=20first=20ternary=20type=20that=20loads=20packed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GGML type 36 (BitNet.cpp's I2_S) now loads as BitNetB158TensorData — the first packed ternary TensorData: 0.25 bytes per weight plus one FP32 scale, instead of the #1033 FP32 widening. Its packedView carries the exact FP32×BITNET_B1_58 dispatch key the ternary f32 kernel pack serves (#1138), so a loaded BitNet weight reaches the vendored NeoGPU LUT kernel with no further conversion. The messy part is the wire format, and it is pinned against the vendored BitNet.cpp sources rather than folklore: the payload's bit order depends on which converter wrote the file — BitNet.cpp packs QK-element blocks high bits first, and QK_I2_S is 128 on the x86 pipeline but 64 on ARM (an architecture-dependent file format); NeoGPU's converter packs sequentially, already the BITNET_B1_58 order. The scale moves too: BitNet.cpp puts one FP32 in a 32-byte trailer after the payload (w = (code-1)*scale); NeoGPU writes a companion _scale tensor defined as divide-by (stored inverted, and consumed rather than delivered as a parameter). The loader takes an i2sLayout knob (GROUP_128 default | GROUP_64 | SEQUENTIAL), repacks once at load (traced as repack-i2s), and accepts either scale source in flavor-preferred order. Byte code 3 fails fast at import — I2sRepack is the only place that validates it; the kernels stay deliberately unvalidating (3 → +2 by LUT arithmetic). GGML_QUANT_SIZES sizes the I2_S payload only; the trailer is read separately from the source, so both flavors' offset math stays truthful. DequantizeTo(FP32) still widens, through the same codec decode. Tests: repack goldens against an independent reimplementation of quantize_i2_s (both QK flavors), code-3 rejection in every layout, wrong-multiple errors naming the other flavor, synthetic-GGUF end-to-end loads for all three flavors (packed data, scale resolution, companion consumption, DequantizeTo parity), and the loader/codec agreement invariant the kernel pack is defined against. Refs #1140, #1136; relates #1033 Co-Authored-By: Claude Fable 5 --- .../kotlin/sk/ainet/io/gguf/Constants.kt | 18 +- .../kotlin/sk/ainet/io/gguf/I2sRepack.kt | 96 ++++++++++ .../io/gguf/StreamingGgufParametersLoader.kt | 134 ++++++++++++++ .../sk/ainet/io/gguf/I2sGgufLoadTest.kt | 165 ++++++++++++++++++ .../kotlin/sk/ainet/io/gguf/I2sRepackTest.kt | 94 ++++++++++ .../lang/tensor/data/BitNetB158TensorData.kt | 106 +++++++++++ .../tensor/data/BitNetB158TensorDataTest.kt | 71 ++++++++ 7 files changed, 683 insertions(+), 1 deletion(-) create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sRepack.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sGgufLoadTest.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sRepackTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorData.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorDataTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/Constants.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/Constants.kt index af302309a..f5afc8202 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/Constants.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/Constants.kt @@ -48,7 +48,18 @@ enum class GGMLQuantizationType(val value: Int) { BF16(30), TQ1_0(34), TQ2_0(35), - // Note: types 31-33 and 36-38 have been removed in llama.cpp + + /** + * BitNet.cpp's ternary type in the slot llama.cpp vacated (types 31-33 and + * 37-38 remain removed there). Not block-regular: the payload is 2-bit + * codes, 4 per byte, and the per-*tensor* FP32 scale sits in a trailer + * after the whole payload (BitNet.cpp) or in a companion `_scale` + * tensor (NeoGPU's converter) — and the payload's bit order additionally + * depends on which converter wrote the file (see `I2sGgufLayout`). The + * loader normalizes all of it to `TensorEncoding.BITNET_B1_58` at load; + * I2_S never becomes a `TensorEncoding` of its own. + */ + I2_S(36), MXFP4(39), /** @@ -113,6 +124,11 @@ val GGML_QUANT_SIZES: Map> = mapOf( GGMLQuantizationType.BF16 to (1 to 2), GGMLQuantizationType.TQ1_0 to (256 to 2 + 4 * 13), GGMLQuantizationType.TQ2_0 to (256 to 2 + 64), + // I2_S: 4 codes per byte. This sizes the PAYLOAD only — the per-tensor + // scale trailer (BitNet.cpp writes 32 bytes after the payload; NeoGPU + // writes none) is deliberately outside the block math, read separately + // by the loader. See GGMLQuantizationType.I2_S. + GGMLQuantizationType.I2_S to (4 to 1), // MXFP4: Microscaling FP4 format - 32 elements per block, 17 bytes (16 for data + 1 for scale) GGMLQuantizationType.MXFP4 to (32 to 17) ) diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sRepack.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sRepack.kt new file mode 100644 index 000000000..a7e4eca5d --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sRepack.kt @@ -0,0 +1,96 @@ +package sk.ainet.io.gguf + +/** + * Which bit order an I2_S GGUF's payload is in — a property of the *converter that wrote the + * file*, not recoverable from the bytes, so the caller has to say (#1140). + * + * BitNet.cpp's `quantize_i2_s` packs 2-bit codes into fixed-size blocks, element `j` of a block + * going to byte `j % (QK/4)`, bit-pair `6 − 2·(j / (QK/4))` (high bits first) — and `QK_I2_S` + * is **128 when the file was quantized on x86 (AVX) and 64 on ARM (NEON)**, an architecture- + * dependent file format. NeoGPU's `convert_bitnet_to_gguf.py` packs the same codes sequentially, + * four consecutive elements per byte, low bit-pair first — byte-identical to SKaiNET's + * `BITNET_B1_58` payload. All three agree on the code mapping `{0,1,2} → {-1,0,+1}`. + */ +public enum class I2sGgufLayout(internal val blockElements: Int) { + /** BitNet.cpp file quantized with the x86/AVX pipeline (`QK_I2_S = 128`, 32-byte blocks). The common case for published GGUFs. */ + GROUP_128(128), + + /** BitNet.cpp file quantized with the ARM/NEON pipeline (`QK_I2_S = 64`, 16-byte blocks). */ + GROUP_64(64), + + /** NeoGPU's converter: sequential 4-per-byte, low bit-pair first — already the `BITNET_B1_58` payload order. */ + SEQUENTIAL(4), +} + +/** + * Repacks an I2_S payload into the sequential `BITNET_B1_58` payload order (#1140). + * + * Byte code 3 is **rejected here, at import** — it has no ternary meaning (the kernels decode it + * as +2 by LUT arithmetic, deliberately unvalidated), so a file that contains it is corrupt and + * the load fails fast instead of silently producing garbage logits. + */ +public object I2sRepack { + + /** + * The sequential `BITNET_B1_58` payload of [elementCount] codes read from [bytes] under + * [layout]. For [I2sGgufLayout.SEQUENTIAL] the payload is validated and copied as-is. + * + * @throws IllegalArgumentException on byte code 3, or when [elementCount] does not fill + * [layout]'s blocks exactly + */ + public fun toSequentialPayload(bytes: ByteArray, elementCount: Int, layout: I2sGgufLayout): ByteArray { + require(elementCount % 4 == 0) { "I2_S element count must be a multiple of 4; got $elementCount" } + val payloadBytes = elementCount / 4 + require(bytes.size >= payloadBytes) { + "I2_S payload needs $payloadBytes bytes for $elementCount elements; got ${bytes.size}" + } + if (layout == I2sGgufLayout.SEQUENTIAL) { + for (i in 0 until payloadBytes) { + val b = bytes[i].toInt() and 0xFF + for (lane in 0 until 4) { + val element = i * 4 + lane + if (element >= elementCount) break + requireValidCode((b shr (lane * 2)) and 3, element) + } + } + return bytes.copyOf(payloadBytes) + } + val qk = layout.blockElements + require(elementCount % qk == 0) { + "a ${layout.name} I2_S tensor must be a multiple of $qk elements; got $elementCount " + + "(the file may be the other BitNet.cpp flavor — try ${otherGroup(layout).name})" + } + val bytesPerBlock = qk / 4 + val out = ByteArray(payloadBytes) + for (element in 0 until elementCount) { + val jb = element % qk + val src = bytes[(element / qk) * bytesPerBlock + jb % bytesPerBlock].toInt() and 0xFF + val code = (src shr (6 - 2 * (jb / bytesPerBlock))) and 3 + requireValidCode(code, element) + val shift = (element % 4) * 2 + out[element / 4] = (out[element / 4].toInt() or (code shl shift)).toByte() + } + return out + } + + /** [payload] with [scale] appended as the little-endian FP32 trailer — a complete `BITNET_B1_58` buffer. */ + public fun withScale(payload: ByteArray, scale: Float): ByteArray { + val out = payload.copyOf(payload.size + 4) + val bits = scale.toRawBits() + out[payload.size] = (bits and 0xFF).toByte() + out[payload.size + 1] = ((bits shr 8) and 0xFF).toByte() + out[payload.size + 2] = ((bits shr 16) and 0xFF).toByte() + out[payload.size + 3] = ((bits shr 24) and 0xFF).toByte() + return out + } + + private fun requireValidCode(code: Int, element: Int) { + require(code != 3) { + "I2_S element $element holds byte code 3, which is not a ternary value — the file is " + + "corrupt, or it was read under the wrong I2sGgufLayout" + } + } + + private fun otherGroup(layout: I2sGgufLayout): I2sGgufLayout = + if (layout == I2sGgufLayout.GROUP_128) I2sGgufLayout.GROUP_64 else I2sGgufLayout.GROUP_128 +} 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 cd8ffef04..c66fe837a 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 @@ -103,6 +103,14 @@ public class StreamingGgufParametersLoader( * Defaults to [NoopTraceSink]: nothing is recorded and nothing is allocated. */ private val traceSink: TraceSink = NoopTraceSink, + /** + * The bit order I2_S (type 36) payloads in this file are in — a property of the converter + * that wrote the file, not recoverable from the bytes (#1140, see [I2sGgufLayout]). Defaults + * to [I2sGgufLayout.GROUP_128], the BitNet.cpp x86 pipeline behind the commonly published + * GGUFs. Wrong-layout loads fail fast on code 3 where possible, but a misdeclared layout can + * also decode silently wrong — this knob is the caller's responsibility. + */ + private val i2sLayout: I2sGgufLayout = I2sGgufLayout.GROUP_128, ) : ParametersLoader { /** @@ -205,6 +213,14 @@ public class StreamingGgufParametersLoader( var current = 0L for (tensorInfo in tensors) { + // NeoGPU's I2_S converter emits a companion `_scale` F32 scalar per ternary + // weight; it is consumed by the I2_S branch below (folded into the BITNET_B1_58 + // trailer), never delivered as a parameter of its own. + if (isI2sCompanionScale(tensorInfo, tensors)) { + current += 1 + onProgress(current, total, tensorInfo.name) + continue + } val tensorForm = formFor(tensorInfo.name) val shape = shapeOf(tensorInfo, tensorForm) // A dense F32 tensor under MAPPED staging never reaches the heap: it is a view over @@ -293,6 +309,11 @@ public class StreamingGgufParametersLoader( GGMLQuantizationType.TQ1_0, GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes, tensorForm) + GGMLQuantizationType.I2_S -> i2sTensor( + ctx, dtype, shape, tensorInfo, rawBytes, tensorForm, + scale = resolveI2sScale(tensorInfo, tensors, reader, source), + ) + else -> throw IllegalStateException( "StreamingGgufParametersLoader: tensor '${tensorInfo.name}' of type " + "${tensorInfo.tensorType} passed the load-time pre-scan but has no load " + @@ -387,6 +408,115 @@ public class StreamingGgufParametersLoader( return ctx.fromData(delivered as sk.ainet.lang.tensor.data.TensorData, dtype) } + /** + * Whether [tensorInfo] is a NeoGPU-converter companion scale — an F32 scalar named + * `_scale` next to an I2_S tensor of that name. Consumed by [resolveI2sScale], + * skipped as a parameter. + */ + private fun isI2sCompanionScale( + tensorInfo: StreamingTensorInfo, + tensors: List, + ): Boolean = + tensorInfo.tensorType == GGMLQuantizationType.F32 && + tensorInfo.name.endsWith("_scale") && + tensors.any { + it.tensorType == GGMLQuantizationType.I2_S && + "${it.name}_scale" == tensorInfo.name + } + + /** + * The per-tensor FP32 scale of an I2_S weight, from wherever its converter put it (#1140): + * + * - **BitNet.cpp** writes it as a trailer after the payload (a 32-byte-aligned region whose + * first 4 bytes are the LE FP32 scale; `w = (code − 1) · scale`). Read directly from the + * source at `absoluteDataOffset + payload` — [StreamingTensorInfo.nBytes] deliberately + * sizes the payload only. + * - **NeoGPU's converter** writes a companion `_scale` F32 scalar, defined as "divide + * the projection output by it" — so the stored multiplier is its inverse. + * - Neither present (or unreadable/non-finite/zero): `1.0`, i.e. the raw codes. Loud in the + * trace via the repack conversion's byte counts, never a crash. + * + * The flavor decides which source is tried first; both are accepted either way, because a + * sequential file with a trailer or a group file with a companion costs nothing to honour. + */ + private fun resolveI2sScale( + tensorInfo: StreamingTensorInfo, + tensors: List, + reader: StreamingGGUFReader, + source: RandomAccessSource, + ): Float { + fun trailer(): Float? = runCatching { + val bytes = source.readAt(tensorInfo.absoluteDataOffset + tensorInfo.nBytes, 4) + val bits = (bytes[0].toInt() and 0xFF) or + ((bytes[1].toInt() and 0xFF) shl 8) or + ((bytes[2].toInt() and 0xFF) shl 16) or + ((bytes[3].toInt() and 0xFF) shl 24) + Float.fromBits(bits) + }.getOrNull()?.takeIf { it.isFinite() && it != 0f } + + fun companionInverse(): Float? { + val companion = tensors.firstOrNull { + it.tensorType == GGMLQuantizationType.F32 && it.name == "${tensorInfo.name}_scale" + } ?: return null + val value = runCatching { bytesToFloatArray(reader.loadTensorData(companion)).firstOrNull() } + .getOrNull() ?: return null + if (!value.isFinite() || value == 0f) return null + return 1f / value + } + + return when (i2sLayout) { + I2sGgufLayout.GROUP_128, I2sGgufLayout.GROUP_64 -> trailer() ?: companionInverse() ?: 1f + I2sGgufLayout.SEQUENTIAL -> companionInverse() ?: trailer() ?: 1f + } + } + + /** + * Materialize an I2_S tensor (#1140): repack the payload into the sequential `BITNET_B1_58` + * order under [i2sLayout] (code 3 fails fast in [I2sRepack]), fold [scale] into the trailer, + * and keep it packed — 0.25 bytes per weight instead of the #1033 FP32 widening. A + * `DequantizeTo` form still gets dense FP32, decoded through the same codec the kernels are + * defined against. + */ + @Suppress("UNCHECKED_CAST") + private fun i2sTensor( + ctx: ExecutionContext, + dtype: KClass, + shape: Shape, + tensorInfo: StreamingTensorInfo, + rawBytes: ByteArray, + tensorForm: WeightForm, + scale: Float, + ): Tensor { + require(dtype == FP32::class || dtype == FP16::class) { + "tensor '${tensorInfo.name}' is I2_S; ternary weights are logically FP32, so the " + + "requested dtype $dtype is not supported" + } + val payload = I2sRepack.toSequentialPayload(rawBytes, tensorInfo.nElements.toInt(), i2sLayout) + val packedBytes = I2sRepack.withScale(payload, scale) + traceConversion( + kind = "repack-i2s", + tensorName = tensorInfo.name, + from = ggufFormat(GGMLQuantizationType.I2_S, rawBytes.size.toLong()), + to = Format(FP32, sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_B1_58), + bytesBefore = rawBytes.size.toLong(), + bytesAfter = packedBytes.size.toLong(), + ) + if (tensorForm.encoding is EncodingRequest.DequantizeTo) { + val dest = sk.ainet.lang.memory.TernaryCodec.decodeBitNet(packedBytes, tensorInfo.nElements.toInt()) + traceConversion( + kind = "widen-i2s", + tensorName = tensorInfo.name, + from = Format(FP32, sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_B1_58), + to = Format.dense(FP32), + bytesBefore = packedBytes.size.toLong(), + bytesAfter = denseFp32Bytes(tensorInfo.nElements), + ) + return ctx.wrapFloatArray(shape, dtype, dest) as Tensor + } + val packed = sk.ainet.lang.tensor.data.BitNetB158TensorData(shape, packedBytes) + return ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + } + /** * [packed] with its blocks permuted into the order the packed matmul kernels read, and *saying * so* (#1120). @@ -489,6 +619,10 @@ public class StreamingGgufParametersLoader( // cannot read the same bytes differently. GGMLQuantizationType.TQ1_0, GGMLQuantizationType.TQ2_0, + // #1140: BitNet.cpp / NeoGPU ternary. Repacked at load into the + // sequential BITNET_B1_58 layout and kept packed (0.25 B/weight) + // — the first ternary type that does NOT widen to FP32 (#1033). + GGMLQuantizationType.I2_S, ) private const val MAX_LISTED_TENSORS = 8 diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sGgufLoadTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sGgufLoadTest.kt new file mode 100644 index 000000000..eb902a2d5 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sGgufLoadTest.kt @@ -0,0 +1,165 @@ +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.tensor.Tensor +import sk.ainet.lang.tensor.data.BitNetB158TensorData +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertTrue + +/** + * #1140 end-to-end: an I2_S GGUF loads **packed** — `BitNetB158TensorData`, 0.25 bytes per + * weight — with the right values, for both converter flavors: + * + * - **BitNet.cpp** (group payload, FP32 scale in a 32-byte trailer after the payload; + * `w = (code−1)·scale`) + * - **NeoGPU** (sequential payload, companion `_scale` F32 scalar defined as "divide the + * output by it") + * + * The first ternary GGUF type that does not take the #1033 FP32 widening. + */ +@OptIn(ExperimentalMemoryApi::class) +class I2sGgufLoadTest { + + private fun randomCodes(count: Int, seed: Int): IntArray { + val rng = Random(seed) + return IntArray(count) { rng.nextInt(3) } + } + + /** BitNet.cpp's `quantize_i2_s` packing + 32-byte trailer whose first 4 bytes are the scale. */ + private fun bitnetCppTensor(name: String, codes: IntArray, qk: Int, scale: Float): SyntheticGguf.TestTensor { + val bytesPerBlock = qk / 4 + val payload = ByteArray(codes.size / 4) + for (j in codes.indices) { + val jb = j % qk + val byteIndex = (j / qk) * bytesPerBlock + jb % bytesPerBlock + payload[byteIndex] = (payload[byteIndex].toInt() or (codes[j] shl (6 - 2 * (jb / bytesPerBlock)))).toByte() + } + val trailer = ByteBuffer.allocate(32).order(ByteOrder.LITTLE_ENDIAN) + repeat(8) { trailer.putFloat(scale) } // BitNet.cpp's converter tiles the scale ×8 + return SyntheticGguf.TestTensor(name, GGMLQuantizationType.I2_S, codes.size.toLong(), payload + trailer.array()) + } + + /** NeoGPU's sequential packing (payload only) + its companion `_scale` scalar. */ + private fun neogpuTensors(name: String, codes: IntArray, weightScale: Float): List { + val payload = ByteArray(codes.size / 4) + for (j in codes.indices) { + payload[j / 4] = (payload[j / 4].toInt() or (codes[j] shl ((j % 4) * 2))).toByte() + } + val scaleBytes = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putFloat(weightScale).array() + return listOf( + SyntheticGguf.TestTensor(name, GGMLQuantizationType.I2_S, codes.size.toLong(), payload), + SyntheticGguf.TestTensor("${name}_scale", GGMLQuantizationType.F32, 1L, scaleBytes), + ) + } + + private fun load( + file: File, + layout: I2sGgufLayout, + form: WeightForm? = null, + ): Map> { + val ctx = DefaultDataExecutionContext() + val loaded = mutableMapOf>() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file) }, + weightForm = form, + i2sLayout = layout, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + @Test + fun bitnetCppFlavorLoadsPackedWithTheTrailerScale() { + val codes = randomCodes(256, seed = 1) + val scale = 0.125f + val file = SyntheticGguf.write(bitnetCppTensor("w", codes, qk = 128, scale = scale)) + try { + val loaded = load(file, I2sGgufLayout.GROUP_128) + assertEquals(setOf("w"), loaded.keys) + val data = assertIs(loaded.getValue("w").data) + assertEquals(scale, data.scale, "trailer scale folded into the BITNET_B1_58 buffer") + assertEquals(256 / 4 + 4, data.packedData.size, "0.25 B/weight + 4 B scale — not widened") + val expected = FloatArray(256) { (codes[it] - 1) * scale } + assertContentEquals(expected, data.toFloatArray(), "values survive the group→sequential repack") + } finally { + file.delete() + } + } + + @Test + fun neogpuFlavorLoadsPackedWithTheInverseCompanionScale() { + val codes = randomCodes(128, seed = 2) + val weightScale = 4.0f // NeoGPU semantics: divide output by it → multiplier 0.25 + val file = SyntheticGguf.write(*neogpuTensors("w", codes, weightScale).toTypedArray()) + try { + val loaded = load(file, I2sGgufLayout.SEQUENTIAL) + assertEquals(setOf("w"), loaded.keys, "the companion _scale tensor is consumed, not delivered") + val data = assertIs(loaded.getValue("w").data) + assertEquals(0.25f, data.scale) + val expected = FloatArray(128) { (codes[it] - 1) * 0.25f } + assertContentEquals(expected, data.toFloatArray()) + } finally { + file.delete() + } + } + + @Test + fun dequantizeToFp32StillWorksAndMatchesThePackedDecode() { + val codes = randomCodes(256, seed = 3) + val file = SyntheticGguf.write(bitnetCppTensor("w", codes, qk = 128, scale = 0.5f)) + try { + val widened = load(file, I2sGgufLayout.GROUP_128, WeightForm(encoding = EncodingRequest.DequantizeTo(FP32))) + val dense = assertIs>(widened.getValue("w").data) + val expected = FloatArray(256) { (codes[it] - 1) * 0.5f } + assertContentEquals(expected, dense.buffer.copyOf(256)) + } finally { + file.delete() + } + } + + @Test + fun group64FlavorRoundTrips() { + val codes = randomCodes(128, seed = 4) + val file = SyntheticGguf.write(bitnetCppTensor("w", codes, qk = 64, scale = 1.0f)) + try { + val data = assertIs(load(file, I2sGgufLayout.GROUP_64).getValue("w").data) + assertContentEquals(FloatArray(128) { (codes[it] - 1).toFloat() }, data.toFloatArray()) + } finally { + file.delete() + } + } + + @Test + fun theLoaderAndTheCodecAgreeOnTheRepackedBytes() { + // The invariant that makes the kernels safe: whatever the loader emits decodes identically + // through TernaryCodec — the reference the ternary kernel pack (#1138) is defined against. + val codes = randomCodes(512, seed = 5) + val file = SyntheticGguf.write(bitnetCppTensor("w", codes, qk = 128, scale = 0.75f)) + try { + val data = assertIs(load(file, I2sGgufLayout.GROUP_128).getValue("w").data) + assertContentEquals( + TernaryCodec.decodeBitNet(data.packedData, 512), + data.toFloatArray(), + ) + assertTrue(data.packedData.size < 512, "packed, not widened") + } finally { + file.delete() + } + } +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sRepackTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sRepackTest.kt new file mode 100644 index 000000000..5f4333ccc --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/I2sRepackTest.kt @@ -0,0 +1,94 @@ +package sk.ainet.io.gguf + +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1140: the group→sequential repack against the packing rule of BitNet.cpp's `quantize_i2_s` + * (element `j` of a `QK`-element block → byte `j % (QK/4)`, bit-pair `6 − 2·(j / (QK/4))`), for + * both of its architecture-dependent flavors, plus the fail-fast on byte code 3. + */ +class I2sRepackTest { + + /** The BitNet.cpp packing, reimplemented independently as the test oracle. */ + private fun packGroup(codes: IntArray, qk: Int): ByteArray { + val bytesPerBlock = qk / 4 + val out = ByteArray(codes.size / 4) + for (j in codes.indices) { + val jb = j % qk + val byteIndex = (j / qk) * bytesPerBlock + jb % bytesPerBlock + out[byteIndex] = (out[byteIndex].toInt() or (codes[j] shl (6 - 2 * (jb / bytesPerBlock)))).toByte() + } + return out + } + + private fun packSequential(codes: IntArray): ByteArray { + val out = ByteArray(codes.size / 4) + for (j in codes.indices) { + out[j / 4] = (out[j / 4].toInt() or (codes[j] shl ((j % 4) * 2))).toByte() + } + return out + } + + private fun randomCodes(count: Int, seed: Int): IntArray { + val rng = Random(seed) + return IntArray(count) { rng.nextInt(3) } // {0, 1, 2} — never 3 + } + + @Test + fun group128RepacksToTheSequentialOrder() { + val codes = randomCodes(256, seed = 1) + val repacked = I2sRepack.toSequentialPayload(packGroup(codes, qk = 128), 256, I2sGgufLayout.GROUP_128) + assertContentEquals(packSequential(codes), repacked) + } + + @Test + fun group64RepacksToTheSequentialOrder() { + val codes = randomCodes(256, seed = 2) + val repacked = I2sRepack.toSequentialPayload(packGroup(codes, qk = 64), 256, I2sGgufLayout.GROUP_64) + assertContentEquals(packSequential(codes), repacked) + } + + @Test + fun sequentialPassesThroughValidatedAndTruncatedToThePayload() { + val codes = randomCodes(64, seed = 3) + val withTrailer = packSequential(codes) + ByteArray(32) { 0x7F } // trailer must be ignored + val out = I2sRepack.toSequentialPayload(withTrailer, 64, I2sGgufLayout.SEQUENTIAL) + assertContentEquals(packSequential(codes), out) + } + + @Test + fun byteCode3FailsFastInEveryLayout() { + for (layout in I2sGgufLayout.entries) { + val bytes = ByteArray(128 / 4) // 128 elements of code 0... + bytes[0] = 0xC0.toByte() // ...except one bit-pair of 3 + val e = assertFailsWith("layout $layout") { + I2sRepack.toSequentialPayload(bytes, 128, layout) + } + assertTrue("code 3" in e.message!!, e.message!!) + } + } + + @Test + fun wrongBlockMultipleNamesTheOtherFlavor() { + // 64 elements fill a GROUP_64 block but not a GROUP_128 one. + val e = assertFailsWith { + I2sRepack.toSequentialPayload(ByteArray(16), 64, I2sGgufLayout.GROUP_128) + } + assertTrue("GROUP_64" in e.message!!, e.message!!) + } + + @Test + fun scaleTrailerIsLittleEndianFp32() { + val buffer = I2sRepack.withScale(ByteArray(4), 0.25f) + val bits = (buffer[4].toInt() and 0xFF) or + ((buffer[5].toInt() and 0xFF) shl 8) or + ((buffer[6].toInt() and 0xFF) shl 16) or + ((buffer[7].toInt() and 0xFF) shl 24) + assertContentEquals(listOf(8), listOf(buffer.size)) + assertTrue(Float.fromBits(bits) == 0.25f) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorData.kt new file mode 100644 index 000000000..660303904 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorData.kt @@ -0,0 +1,106 @@ +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 ternary tensor data in the [TensorEncoding.BITNET_B1_58] layout — the first packed + * ternary `TensorData` (#1040/#1140), closing the "#1033 widens everything to FP32" gap for the + * per-tensor BitNet encoding. + * + * Buffer layout, exactly what [TernaryCodec.encodeBitNet] writes and the ternary kernels read: + * `ceil(volume / 4)` payload bytes (2-bit codes, four consecutive elements per byte, low bit-pair + * first, code `{0,1,2} → {-1,0,+1}`) followed by one little-endian FP32 per-tensor [scale]. + * + * The whole tensor is one block ([blockCount] `== 1`): the encoding has a single scale, so there + * is no per-block structure to expose, and [packedView] carries + * `Format(FP32, BITNET_B1_58) × BLOCKED_ROW_MAJOR` — the exact key the ternary f32 kernel pack + * serves (#1138). `get` returns the *signed code* (−1, 0, +1; byte code 3 → +2), scale not + * applied; decoding with the scale goes through [dequantizeBlock] / [PackedBlockStorage.toFloatArray]. + */ +public class BitNetB158TensorData( + 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()) + private val strides: IntArray = shape.computeStrides() + + override val encoding: TensorEncoding get() = TensorEncoding.BITNET_B1_58 + override val blockCount: Int get() = 1 + override val blockSize: Int get() = shape.volume + override val packedData: ByteArray get() = data + + /** The per-tensor FP32 scale (the trailing 4 bytes). */ + public val scale: Float get() = TernaryCodec.bitNetScale(data, shape.volume) + + init { + val required = (TensorEncoding.BITNET_B1_58.physicalBytes(shape.volume.toLong()) + ?: error("BITNET_B1_58 cannot size ${shape.volume} elements")) + require(data.size >= required) { + "BitNetB158TensorData: buffer is ${data.size} bytes, need >= $required " + + "(ceil(${shape.volume}/4) payload + 4-byte FP32 scale)" + } + } + + override fun dequantizeBlock(blockIdx: Int, output: FloatArray, outputOffset: Int) { + require(blockIdx == 0) { "BITNET_B1_58 is per-tensor: only block 0 exists, got $blockIdx" } + val s = scale + for (i in 0 until shape.volume) { + val outIdx = outputOffset + i + if (outIdx >= output.size) return + output[outIdx] = codeAt(i) * s + } + } + + /** The signed ternary code of flat element [flatIndex] — scale not applied. */ + private fun codeAt(flatIndex: Int): Int = + (((data[flatIndex / 4].toInt() and 0xFF) shr ((flatIndex % 4) * 2)) and 3) - 1 + + override fun get(vararg indices: Int): Byte = codeAt(calcFlatIndex(indices)).toByte() + + override fun set(vararg indices: Int, value: Byte) { + require(value in -1..1) { "BITNET_B1_58 stores ternary codes; got $value" } + val flatIndex = calcFlatIndex(indices) + val byteIndex = flatIndex / 4 + val shift = (flatIndex % 4) * 2 + val cleared = data[byteIndex].toInt() and (3 shl shift).inv() + data[byteIndex] = (cleared or ((value + 1) shl shift)).toByte() + } + + private fun calcFlatIndex(indices: IntArray): Int { + require(indices.size == shape.dimensions.size) { + "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})" + } + var flatIndex = 0 + for (i in indices.indices) { + val idx = indices[i] + require(idx >= 0 && idx < shape.dimensions[i]) { + "Index $idx out of bounds for dimension $i with size ${shape.dimensions[i]}" + } + flatIndex += idx * strides[i] + } + return flatIndex + } + + public companion object { + /** Wrap raw `payload + scale` bytes (validates the size). */ + public fun fromRawBytes(shape: Shape, bytes: ByteArray): BitNetB158TensorData = + BitNetB158TensorData(shape, bytes) + + /** Encode [values] with [TernaryCodec.encodeBitNet] (absmean ternarization). */ + public fun fromFloats(shape: Shape, values: FloatArray): BitNetB158TensorData { + require(values.size == shape.volume) { + "values (${values.size}) must match shape volume (${shape.volume})" + } + return BitNetB158TensorData(shape, TernaryCodec.encodeBitNet(values)) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorDataTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorDataTest.kt new file mode 100644 index 000000000..6f590669c --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/BitNetB158TensorDataTest.kt @@ -0,0 +1,71 @@ +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.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith + +/** + * #1140: the first packed ternary `TensorData`. Its decode must be pinned to [TernaryCodec] — + * the same reference the kernels and the GGUF loader are defined against — so all three readers + * agree about the same bytes. + */ +@OptIn(ExperimentalMemoryApi::class) +class BitNetB158TensorDataTest { + + private fun ternaryValues(count: Int, seed: Int): FloatArray { + var s = seed + return FloatArray(count) { + s = s * 1103515245 + 12345 + ((s ushr 16) % 3 - 1) * 0.5f + } + } + + @Test + fun roundTripsThroughTheCodec() { + val n = 6; val k = 32 + val values = ternaryValues(n * k, seed = 5) + val data = BitNetB158TensorData.fromFloats(Shape(n, k), values) + + assertEquals(TernaryCodec.bitNetScale(data.packedData, n * k), data.scale) + assertContentEquals( + TernaryCodec.decodeBitNet(data.packedData, n * k), + data.toFloatArray(), + "toFloatArray must equal the codec's decode of the same bytes", + ) + } + + @Test + fun getReturnsSignedCodesAndSetWritesThem() { + val data = BitNetB158TensorData.fromFloats(Shape(8), floatArrayOf(1f, -1f, 0f, 1f, -1f, 0f, 0f, 1f)) + assertEquals(1, data.get(0).toInt()) + assertEquals(-1, data.get(1).toInt()) + assertEquals(0, data.get(2).toInt()) + data.set(2, value = -1) + assertEquals(-1, data.get(2).toInt()) + assertFailsWith { data.set(0, value = 2) } + } + + @Test + fun viewCarriesTheExactDispatchFormat() { + val data = BitNetB158TensorData.fromFloats(Shape(4, 8), ternaryValues(32, seed = 9)) + val view = data.packedView + assertEquals(TensorEncoding.BITNET_B1_58, view.format.encoding) + assertEquals(1, data.blockCount, "per-tensor encoding: one block") + assertEquals(32, data.blockSize) + // the view decodes with the scale applied, exactly like the codec + val decoded = TernaryCodec.decodeBitNet(data.packedData, 32) + assertEquals(decoded[9], view.get(1, 1)) + } + + @Test + fun rejectsBuffersWithoutTheScaleTrailer() { + assertFailsWith { + BitNetB158TensorData(Shape(8), ByteArray(2)) // payload only, no trailer + } + } +}