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 + } + } +}