From 69dd04ef19949532396e5c7e24e59201ea8f9be3 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sat, 29 Aug 2026 16:38:01 +0200 Subject: [PATCH] feat(io): AOT GGUF converter for I2_S -- the eager-exec half of #1207 The on-device sidecar cache (#1204) pays its "convert once" cost by doubling the first cold start on the constrained device this effort exists to protect. An AOT converter that runs once, offline, ahead of the device ever seeing the file, avoids that entirely -- a build that controls its own model pipeline gets a genuinely zero-runtime-cost load, no sidecar, no first-load penalty. - GgufTensorEntry gains an optional `rawBytes: ByteArray?` alternative to `tensor`: a converter that already has the exact on-disk bytes for an 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 GGML_QUANT_SIZES block math -- an I2_S BITNET_B1_58 buffer with its trailing FP32 scale) can hand them over directly. GGUFWriter's expectedTensorSize/materializeTensor use rawBytes.size and the bytes verbatim when present, bypassing TensorFlatten (which is element-indexed and can't represent an opaque quantized blob) and GGML_QUANT_SIZES entirely. Every existing tensor-mode caller/test is unaffected. - StreamingGgufParametersLoader's I2_S scale-resolution logic (resolveI2sScale, i2sCompanionScaleTensor, i2sTrailerScale, i2sTrailerScaleIsMappable) is extracted into a shared I2sScale.kt so the new converter gets the identical answer the streaming loader would, rather than a second, potentially-drifting copy. - I2sAotConverter.convert(): reads an arbitrary GGUF via StreamingGGUFReader, repacks I2_S tensors into SEQUENTIAL+trailer order (I2sRepack, the same repack the loader already does, just ahead of time), drops the now-redundant companion scale tensor (keeping one would defeat #1203's mmap fast path, which requires no companion to exist), and passes every other tensor and all KV metadata through unchanged. Produces a GgufWriteRequest for GGUFWriter to write out. New I2sAotConverterTest: converts a GROUP_128 (BitNet.cpp) + companion-free trailer source file, asserts the converted I2_S tensor decodes identically to the source AND actually takes #1203's mmap fast path afterwards (packedStorage is Storage.OffHeap under WeightResidency.MAPPED) -- proving this is a real substitute for the on-device cache, not just the same cost moved earlier. A passthrough Q4_0 tensor is asserted unchanged by the round trip. Scope: this is the eager-exec leg only. The IREE/.irpa leg (also named in #1207) is not addressed here. --- .../kotlin/sk/ainet/io/gguf/I2sScale.kt | 91 +++++++++++++ .../io/gguf/StreamingGgufParametersLoader.kt | 83 +----------- .../sk/ainet/io/gguf/export/GGUFWriter.kt | 65 +++++---- .../ainet/io/gguf/export/GgufExportFacade.kt | 28 +++- .../ainet/io/gguf/export/I2sAotConverter.kt | 86 ++++++++++++ .../io/gguf/export/I2sAotConverterTest.kt | 124 ++++++++++++++++++ 6 files changed, 363 insertions(+), 114 deletions(-) create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/I2sScale.kt create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/export/I2sAotConverter.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/export/I2sAotConverterTest.kt 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() + } + } +}