diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sScale.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sScale.kt new file mode 100644 index 000000000..d92461389 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sScale.kt @@ -0,0 +1,91 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.RandomAccessSource + +/** + * I2_S scale resolution (#1140), shared by [StreamingGgufParametersLoader] and any other reader + * of an I2_S GGUF — notably an AOT converter (#1207) that needs the exact same answer the + * streaming loader would give, without duplicating the decision. + */ + +/** The `_scale` companion tensor for an I2_S weight, if the converter wrote one (#1140). */ +internal fun i2sCompanionScaleTensor( + tensorInfo: StreamingTensorInfo, + tensors: List, +): StreamingTensorInfo? = tensors.firstOrNull { + it.tensorType == GGMLQuantizationType.F32 && it.name == "${tensorInfo.name}_scale" +} + +private fun le32(bytes: ByteArray, offset: Int = 0): 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 little-endian FP32 immediately after [tensorInfo]'s payload — BitNet.cpp's trailer + * convention — or `null` if unreadable, non-finite, or zero. [StreamingTensorInfo.nBytes] + * deliberately sizes the payload only, so `absoluteDataOffset + nBytes` is exactly where a + * trailer would start. + */ +internal fun i2sTrailerScale(tensorInfo: StreamingTensorInfo, source: RandomAccessSource): Float? = + runCatching { le32(source.readAt(tensorInfo.absoluteDataOffset + tensorInfo.nBytes, 4)) } + .getOrNull()?.takeIf { it.isFinite() && it != 0f } + +/** + * 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. + * + * [layout] 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. + */ +internal fun resolveI2sScale( + tensorInfo: StreamingTensorInfo, + tensors: List, + reader: StreamingGGUFReader, + source: RandomAccessSource, + layout: I2sGgufLayout, +): Float { + fun companionInverse(): Float? { + val companion = i2sCompanionScaleTensor(tensorInfo, tensors) ?: return null + val value = runCatching { le32(reader.loadTensorData(companion)) }.getOrNull() ?: return null + if (!value.isFinite() || value == 0f) return null + return 1f / value + } + + return when (layout) { + I2sGgufLayout.GROUP_128, I2sGgufLayout.GROUP_64 -> i2sTrailerScale(tensorInfo, source) ?: companionInverse() ?: 1f + I2sGgufLayout.SEQUENTIAL -> companionInverse() ?: i2sTrailerScale(tensorInfo, source) ?: 1f + } +} + +/** + * Whether an I2_S tensor's on-disk bytes are, as-is, a complete kernel-ready `BITNET_B1_58` + * buffer — payload immediately followed by its own trailing FP32 scale — so mapping + * `[absoluteDataOffset, absoluteDataOffset + nBytes + 4)` directly gives exactly what + * [resolveI2sScale] would have computed anyway (#1203). + * + * Only ever true for [I2sGgufLayout.SEQUENTIAL]: the payload itself is already in + * `BITNET_B1_58` order there, whereas `GROUP_128`/`GROUP_64` payloads still need permuting + * regardless of where the scale lives. False whenever a companion `_scale` tensor exists + * — [resolveI2sScale]'s `SEQUENTIAL` order prefers it over a trailer — or the trailer bytes + * don't parse to a finite, nonzero float. + */ +internal fun i2sTrailerScaleIsMappable( + tensorInfo: StreamingTensorInfo, + tensors: List, + source: RandomAccessSource, + layout: I2sGgufLayout, +): Boolean = + layout == I2sGgufLayout.SEQUENTIAL && + i2sCompanionScaleTensor(tensorInfo, tensors) == null && + i2sTrailerScale(tensorInfo, source) != null 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 68a74d0bc..803f4cf7e 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 @@ -317,7 +317,7 @@ public class StreamingGgufParametersLoader( // tensor overrides them) -- otherwise fall through to i2sTensor's // repack, unchanged. GGMLQuantizationType.I2_S -> - if (i2sTrailerScaleIsMappable(tensorInfo, tensors, source)) { + if (i2sTrailerScaleIsMappable(tensorInfo, tensors, source, i2sLayout)) { sk.ainet.lang.tensor.storage.TensorEncoding.BITNET_B1_58 } else { null @@ -420,7 +420,7 @@ public class StreamingGgufParametersLoader( GGMLQuantizationType.I2_S -> i2sTensor( ctx, dtype, shape, tensorInfo, rawBytes, tensorForm, - scale = resolveI2sScale(tensorInfo, tensors, reader, source), + scale = resolveI2sScale(tensorInfo, tensors, reader, source, i2sLayout), ) else -> throw IllegalStateException( @@ -540,85 +540,6 @@ public class StreamingGgufParametersLoader( "${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 companionInverse(): Float? { - val companion = i2sCompanionScaleTensor(tensorInfo, tensors) ?: 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 -> i2sTrailerScale(tensorInfo, source) ?: companionInverse() ?: 1f - I2sGgufLayout.SEQUENTIAL -> companionInverse() ?: i2sTrailerScale(tensorInfo, source) ?: 1f - } - } - - /** The `_scale` companion tensor for an I2_S weight, if the converter wrote one (#1140). */ - private fun i2sCompanionScaleTensor( - tensorInfo: StreamingTensorInfo, - tensors: List, - ): StreamingTensorInfo? = tensors.firstOrNull { - it.tensorType == GGMLQuantizationType.F32 && it.name == "${tensorInfo.name}_scale" - } - - /** - * The little-endian FP32 immediately after [tensorInfo]'s payload — BitNet.cpp's trailer - * convention — or `null` if unreadable, non-finite, or zero. [StreamingTensorInfo.nBytes] - * deliberately sizes the payload only, so `absoluteDataOffset + nBytes` is exactly where a - * trailer would start. - */ - private fun i2sTrailerScale(tensorInfo: StreamingTensorInfo, source: RandomAccessSource): 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 } - - /** - * Whether an I2_S tensor's on-disk bytes are, as-is, a complete kernel-ready `BITNET_B1_58` - * buffer — payload immediately followed by its own trailing FP32 scale — so mapping - * `[absoluteDataOffset, absoluteDataOffset + nBytes + 4)` directly gives exactly what - * [resolveI2sScale] would have computed anyway (#1203). - * - * Only ever true for [I2sGgufLayout.SEQUENTIAL]: the payload itself is already in - * `BITNET_B1_58` order there, whereas `GROUP_128`/`GROUP_64` payloads still need permuting - * regardless of where the scale lives. False whenever a companion `_scale` tensor - * exists — [resolveI2sScale]'s `SEQUENTIAL` order prefers it over a trailer — or the trailer - * bytes don't parse to a finite, nonzero float. - */ - private fun i2sTrailerScaleIsMappable( - tensorInfo: StreamingTensorInfo, - tensors: List, - source: RandomAccessSource, - ): Boolean = - i2sLayout == I2sGgufLayout.SEQUENTIAL && - i2sCompanionScaleTensor(tensorInfo, tensors) == null && - i2sTrailerScale(tensorInfo, source) != null - /** * 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, diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GGUFWriter.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GGUFWriter.kt index ffc28e4e0..21a8821c8 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GGUFWriter.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GGUFWriter.kt @@ -125,7 +125,8 @@ public object GGUFWriter { require(entry.shape.all { it > 0 }) { "Tensor ${entry.ggufName} has non-positive dimensions ${entry.shape}" } - require(GGML_QUANT_SIZES.containsKey(entry.quantization)) { + // A rawBytes entry declares its own size and needs no block-size metadata at all. + require(entry.rawBytes != null || GGML_QUANT_SIZES.containsKey(entry.quantization)) { "Quantization ${entry.quantization} missing size metadata" } } @@ -209,16 +210,21 @@ public object GGUFWriter { } private fun materializeTensor(entry: GgufTensorEntry, expectedSize: Int): ByteArray { - val bytes = when (entry.quantization) { - GGMLQuantizationType.F32 -> materializeF32(entry) - GGMLQuantizationType.F16 -> materializeF16(entry) - GGMLQuantizationType.BF16 -> materializeBF16(entry) - GGMLQuantizationType.F64 -> materializeF64(entry) - GGMLQuantizationType.I8 -> materializeI8(entry) - GGMLQuantizationType.I16 -> materializeI16(entry) - GGMLQuantizationType.I32 -> materializeI32(entry) - GGMLQuantizationType.I64 -> materializeI64(entry) - else -> materializeRaw(entry) + val bytes = entry.rawBytes ?: run { + val tensor = checkNotNull(entry.tensor) { + "GgufTensorEntry '${entry.ggufName}' has neither tensor nor rawBytes" + } + when (entry.quantization) { + GGMLQuantizationType.F32 -> materializeF32(tensor) + GGMLQuantizationType.F16 -> materializeF16(tensor) + GGMLQuantizationType.BF16 -> materializeBF16(tensor) + GGMLQuantizationType.F64 -> materializeF64(tensor) + GGMLQuantizationType.I8 -> materializeI8(tensor) + GGMLQuantizationType.I16 -> materializeI16(tensor) + GGMLQuantizationType.I32 -> materializeI32(tensor) + GGMLQuantizationType.I64 -> materializeI64(tensor) + else -> materializeRaw(tensor) + } } require(bytes.size == expectedSize) { "Tensor ${entry.ggufName} size mismatch: expected $expectedSize, got ${bytes.size}" @@ -226,65 +232,66 @@ public object GGUFWriter { return bytes } - private fun materializeF32(entry: GgufTensorEntry): ByteArray { - val floatData = TensorFlatten.flattenFloats(entry.tensor) + private fun materializeF32(tensor: Tensor<*, *>): ByteArray { + val floatData = TensorFlatten.flattenFloats(tensor) val out = ByteWriter() floatData.forEach { out.writeFloat32(it) } return out.toByteArray() } - private fun materializeF16(entry: GgufTensorEntry): ByteArray { - val floatData = TensorFlatten.flattenFloats(entry.tensor) + private fun materializeF16(tensor: Tensor<*, *>): ByteArray { + val floatData = TensorFlatten.flattenFloats(tensor) val out = ByteWriter() floatData.forEach { out.writeUInt16(floatToHalfBits(it).toUShort()) } return out.toByteArray() } - private fun materializeBF16(entry: GgufTensorEntry): ByteArray { - val floatData = TensorFlatten.flattenFloats(entry.tensor) + private fun materializeBF16(tensor: Tensor<*, *>): ByteArray { + val floatData = TensorFlatten.flattenFloats(tensor) val out = ByteWriter() floatData.forEach { out.writeUInt16(bfloat16Bits(it).toUShort()) } return out.toByteArray() } - private fun materializeF64(entry: GgufTensorEntry): ByteArray { - val doubleData = TensorFlatten.flattenDoubles(entry.tensor) + private fun materializeF64(tensor: Tensor<*, *>): ByteArray { + val doubleData = TensorFlatten.flattenDoubles(tensor) val out = ByteWriter() doubleData.forEach { out.writeFloat64(it) } return out.toByteArray() } - private fun materializeI8(entry: GgufTensorEntry): ByteArray { - val bytes = TensorFlatten.flattenBytes(entry.tensor) + private fun materializeI8(tensor: Tensor<*, *>): ByteArray { + val bytes = TensorFlatten.flattenBytes(tensor) return bytes } - private fun materializeI16(entry: GgufTensorEntry): ByteArray { - val values = TensorFlatten.flattenShorts(entry.tensor) + private fun materializeI16(tensor: Tensor<*, *>): ByteArray { + val values = TensorFlatten.flattenShorts(tensor) val out = ByteWriter() values.forEach { out.writeUInt16(it.toUShort()) } return out.toByteArray() } - private fun materializeI32(entry: GgufTensorEntry): ByteArray { - val ints = TensorFlatten.flattenInts(entry.tensor) + private fun materializeI32(tensor: Tensor<*, *>): ByteArray { + val ints = TensorFlatten.flattenInts(tensor) val out = ByteWriter() ints.forEach { out.writeInt32(it) } return out.toByteArray() } - private fun materializeI64(entry: GgufTensorEntry): ByteArray { - val longs = TensorFlatten.flattenLongs(entry.tensor) + private fun materializeI64(tensor: Tensor<*, *>): ByteArray { + val longs = TensorFlatten.flattenLongs(tensor) val out = ByteWriter() longs.forEach { out.writeInt64(it) } return out.toByteArray() } - private fun materializeRaw(entry: GgufTensorEntry): ByteArray { - return TensorFlatten.flattenBytes(entry.tensor) + private fun materializeRaw(tensor: Tensor<*, *>): ByteArray { + return TensorFlatten.flattenBytes(tensor) } private fun expectedTensorSize(entry: GgufTensorEntry): Int { + entry.rawBytes?.let { return it.size } val (blockSize, typeSize) = GGML_QUANT_SIZES[entry.quantization] ?: error("Quantization ${entry.quantization} missing size metadata") val volume = entry.shape.fold(1L) { acc, d -> acc * max(1, d).toLong() } diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GgufExportFacade.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GgufExportFacade.kt index 0156ec26b..4d6e05110 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GgufExportFacade.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/GgufExportFacade.kt @@ -36,13 +36,33 @@ public data class GgufExportOptions( val provenance: Map = emptyMap() ) -/** Tensor entry to be consumed by a future GGUF writer implementation. */ +/** + * Tensor entry to be consumed by [sk.ainet.io.gguf.export.GGUFWriter]. + * + * Exactly one of [tensor] or [rawBytes] must be set. [tensor] is the original path: a SKaiNET + * tensor whose elements the writer encodes itself (dense float widths, or [rawBytes]-free raw + * passthrough via `TensorFlatten.flattenBytes`, which needs a real element-indexed `Tensor`). + * + * [rawBytes] (#1207) is for a converter that already has the exact on-disk bytes for this entry + * — a passthrough of a quantized blob this writer doesn't need to interpret, or a buffer whose + * true size exceeds its type's formal [sk.ainet.io.gguf.GGML_QUANT_SIZES] block math entirely + * (e.g. an I2_S `BITNET_B1_58` buffer with its trailing FP32 scale). When set, [tensor] is + * ignored, [TensorFlatten] is bypassed, and the entry's written size is exactly + * `rawBytes.size` — not derived from [quantization]/[shape] at all. + */ public data class GgufTensorEntry( val ggufName: String, - val tensor: Tensor<*, *>, + val tensor: Tensor<*, *>? = null, val quantization: GGMLQuantizationType, - val shape: List -) + val shape: List, + val rawBytes: ByteArray? = null, +) { + init { + require((tensor == null) != (rawBytes == null)) { + "GgufTensorEntry '$ggufName' needs exactly one of tensor or rawBytes" + } + } +} /** Aggregate export payload prepared by the facade; writer will consume this. */ public data class GgufWriteRequest( diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/I2sAotConverter.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/I2sAotConverter.kt new file mode 100644 index 000000000..d9200a2f2 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/I2sAotConverter.kt @@ -0,0 +1,86 @@ +package sk.ainet.io.gguf.export + +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.gguf.GGMLQuantizationType +import sk.ainet.io.gguf.I2sGgufLayout +import sk.ainet.io.gguf.I2sRepack +import sk.ainet.io.gguf.StreamingGGUFReader +import sk.ainet.io.gguf.StreamingTensorInfo +import sk.ainet.io.gguf.i2sCompanionScaleTensor +import sk.ainet.io.gguf.resolveI2sScale + +/** + * AOT conversion (#1207): reads an arbitrary GGUF file and re-encodes any I2_S (ternary/BitNet) + * tensors into a `SEQUENTIAL`, trailer-scaled buffer — the exact shape #1203's loader serves + * zero-copy from mmap, with no runtime repack at all. Every other tensor, and every file-level + * KV metadata entry, passes through byte-for-byte/value-for-value unchanged. + * + * This is the "convert once, ahead of time" alternative to caching the repack on-device + * (#1204): a build that controls its own model pipeline runs this once, ships the converted + * file, and every future load of it is a genuine zero-copy mmap — no sidecar, no first-load + * penalty on the constrained device the conversion exists to protect. See #1198 for the full + * reasoning behind preferring this over the on-device cache. + * + * Scope: this converts the file *content* a consuming loader reads (tensors + KV metadata). It + * does not attempt to reconstruct duplicate top-level KV keys — a pathological, practically + * unseen case `StreamingGGUFReader` itself only handles defensively (renaming them + * `_dup_N`) — those would round-trip as literal `_dup_N`-suffixed keys rather than true + * duplicates. + */ +public object I2sAotConverter { + + /** Reader-synthetic keys ([StreamingGGUFReader.fields] adds these; they are not real GGUF KV entries — [GGUFWriter] computes its own header counts/version). */ + private val SYNTHETIC_KEYS = setOf("GGUF.tensor_count", "GGUF.kv_count", "GGUF.version") + + /** + * Build the [GgufWriteRequest] for converting [source]'s tensors, whose I2_S tensors (if + * any) are in [sourceLayout] order (see [I2sGgufLayout] — a property of the converter that + * wrote the *source* file, not recoverable from its bytes). The caller writes the result + * with [GGUFWriter], e.g. `GGUFWriter.writeToSink(convert(source, layout), sink)`. + */ + public fun convert(source: RandomAccessSource, sourceLayout: I2sGgufLayout): GgufWriteRequest { + StreamingGGUFReader.open(source).use { reader -> + val tensors = reader.tensors + val entries = ArrayList(tensors.size) + for (info in tensors) { + when { + info.tensorType == GGMLQuantizationType.I2_S -> { + val scale = resolveI2sScale(info, tensors, reader, source, sourceLayout) + val sequential = I2sRepack.toSequentialPayload( + reader.loadTensorData(info), + info.nElements.toInt(), + sourceLayout, + ) + entries += GgufTensorEntry( + ggufName = info.name, + quantization = GGMLQuantizationType.I2_S, + shape = info.shape.map { it.toInt() }, + rawBytes = I2sRepack.withScale(sequential, scale), + ) + } + // Its value is now folded into the trailer above — the converted I2_S + // tensor is self-contained, so the companion scalar is no longer needed + // (and keeping it would defeat #1203's mmap fast path, which requires no + // companion tensor to exist at all). + isI2sCompanionScale(info, tensors) -> Unit + else -> entries += GgufTensorEntry( + ggufName = info.name, + quantization = info.tensorType, + shape = info.shape.map { it.toInt() }, + rawBytes = reader.loadTensorData(info), + ) + } + } + val metadata = LinkedHashMap(reader.fields.size) + for ((key, value) in reader.fields) { + if (key in SYNTHETIC_KEYS || value == null) continue + metadata[key] = value + } + return GgufWriteRequest(metadata = metadata, tensors = entries, tensorMap = emptyMap()) + } + } + + private fun isI2sCompanionScale(tensorInfo: StreamingTensorInfo, tensors: List): Boolean = + tensorInfo.tensorType == GGMLQuantizationType.F32 && + tensors.any { it.tensorType == GGMLQuantizationType.I2_S && i2sCompanionScaleTensor(it, tensors) == tensorInfo } +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/export/I2sAotConverterTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/export/I2sAotConverterTest.kt new file mode 100644 index 000000000..722b433b7 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/export/I2sAotConverterTest.kt @@ -0,0 +1,124 @@ +package sk.ainet.io.gguf.export + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.gguf.GGMLQuantizationType +import sk.ainet.io.gguf.I2sGgufLayout +import sk.ainet.io.gguf.StreamingGgufParametersLoader +import sk.ainet.io.gguf.SyntheticGguf +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightResidency +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.BitNetB158TensorData +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertIs + +/** + * #1207: [I2sAotConverter] converts a `GROUP_128` (BitNet.cpp) I2_S tensor into a + * `SEQUENTIAL`, trailer-scaled one — the shape #1203's loader serves as true zero-copy mmap — + * while a non-I2_S tensor and the file's KV metadata pass through unchanged. The converted + * file's decoded values must equal the source file's, and it must actually take the mmap fast + * path afterwards, proving this is a real substitute for #1204's on-device sidecar cache, not + * just a same-cost repack moved earlier. + */ +@OptIn(ExperimentalMemoryApi::class) +class I2sAotConverterTest { + + private fun groupPayload(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 leFloat(value: Float): ByteArray = + ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putFloat(value).array() + + private fun randomCodes(count: Int, seed: Int): IntArray { + val rng = Random(seed) + return IntArray(count) { rng.nextInt(3) } // {0, 1, 2} — never 3 + } + + private fun sourceFile(): File { + val elements = 256 + val codes = randomCodes(elements, seed = 9) + val i2s = groupPayload(codes, qk = 128) + leFloat(0.5f) // BitNet.cpp GROUP_128 + trailer + return SyntheticGguf.write( + SyntheticGguf.TestTensor("attn.w", GGMLQuantizationType.I2_S, elements.toLong(), i2s), + SyntheticGguf.tensor("attn.q4_0", GGMLQuantizationType.Q4_0, elements = 256, seed = 3), + ) + } + + private fun convertedFile(source: File): File { + val request = I2sAotConverter.convert( + source = JvmRandomAccessSource.open(source), + sourceLayout = I2sGgufLayout.GROUP_128, + ) + val (_, bytes) = GGUFWriter.writeToByteArray(request) + val out = Files.createTempFile("skainet-aot-converted", ".gguf") + Files.write(out, bytes) + return out.toFile() + } + + private fun load(f: File, layout: I2sGgufLayout, form: WeightForm): Map> { + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightForm = form, + i2sLayout = layout, + ).load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + @Test + fun convertedFileDecodesIdenticallyAndTakesTheMmapFastPath() { + val source = sourceFile() + try { + val converted = convertedFile(source) + try { + val fromSource = load(source, I2sGgufLayout.GROUP_128, WeightForm(residency = WeightResidency.HEAP)) + val fromConverted = load(converted, I2sGgufLayout.SEQUENTIAL, WeightForm(residency = WeightResidency.MAPPED)) + + val sourceI2s = assertIs(fromSource.getValue("attn.w").data) + val convertedI2s = assertIs(fromConverted.getValue("attn.w").data) + assertContentEquals( + sourceI2s.toFloatArray(), convertedI2s.toFloatArray(), + "converted I2_S values must equal the source's", + ) + assertIs( + convertedI2s.packedStorage, + "the converted file's I2_S tensor must take the mmap fast path — no repack, no heap copy", + ) + + val sourceQ4 = fromSource.getValue("attn.q4_0").data as PackedBlockStorage + val convertedQ4 = fromConverted.getValue("attn.q4_0").data as PackedBlockStorage + assertContentEquals( + sourceQ4.toFloatArray(), convertedQ4.toFloatArray(), + "a passthrough (non-I2_S) tensor must be unchanged by conversion", + ) + } finally { + converted.delete() + } + } finally { + source.delete() + } + } +}