From 0ed276d12ee99bcd792f34ee28513abacb7b8f81 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Sun, 25 Oct 2026 00:59:00 +0200 Subject: [PATCH] gguf: parse metadata through a buffered cursor StreamingGGUFReader read every metadata value with its own RandomAccessSource.readAt (a fresh byte array each, one FileChannel read each on the file sources) and allocated a Pair and a boxed offset per array element. A tokenizer GGUF carries arrays of ~262k strings, floats and ints, so opening it cost over a million source calls. GgufMetadataCursor reads the source in 64 KB blocks and decodes little-endian values from the block; string and 32-bit arrays are decoded block by block straight into pre-sized lists. `fields` keeps its keys, order, value classes and element classes, and stays plain lists (a lazily boxing view was slower for consumers that walk the lists several times on an interpreted runtime). Truncated input and array counts larger than the file fail with IllegalArgumentException before allocating. StreamingGGUFReaderMetadataParityTest compares against the previous per-value parse (kept in the test) on every value type, a long string over the block size, a 262k-entry tokenizer-shaped file, a source that returns 7 bytes per read, the file source, truncated input and a GGUFWriter round trip; SKAINET_GGUF_FILE adds a real file. GgufTokenizerLoadBenchmarkTest times parse and tokenizer build. Gemma-3 tokenizer GGUF, JVM: file source 1,002 -> 21-45 ms, in-memory 55 -> 21-34 ms. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 + .../GgufTokenizerLoadBenchmarkTest.kt | 56 ++++ .../sk/ainet/io/gguf/GgufMetadataCursor.kt | 177 +++++++++++ .../sk/ainet/io/gguf/StreamingGGUFReader.kt | 236 +++++--------- .../StreamingGGUFReaderMetadataParityTest.kt | 290 ++++++++++++++++++ 5 files changed, 599 insertions(+), 164 deletions(-) create mode 100644 skainet-io/skainet-io-core/src/jvmTest/kotlin/sk/ainet/io/tokenizer/GgufTokenizerLoadBenchmarkTest.kt create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMetadataCursor.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StreamingGGUFReaderMetadataParityTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ea3a60ba..b3b040453 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- **`StreamingGGUFReader` parses the metadata through a buffered cursor** (`skainet-io-gguf`): the header, KV and tensor-info sections used one `RandomAccessSource.readAt` per value (a fresh byte array each, and on `JvmRandomAccessSource`/`AndroidRandomAccessSource` one `FileChannel` read each), plus a `Pair` and a boxed offset per array element. A tokenizer GGUF has arrays of ~262k strings, floats and ints, so that was over a million source calls. The reader now pulls the source in 64 KB blocks and decodes values from the block; the large arrays are decoded block by block straight into pre-sized lists. `fields` is unchanged: same keys and order, same value and element classes (`UInt` for UINT32, `Float` for FLOAT32, …), still plain lists. A truncated file or an array count larger than the file now fails with an `IllegalArgumentException` before allocating. Gemma-3 tokenizer GGUF (262k tokens), JVM: file source 1,002 → 21–45 ms, in-memory 55 → 21–34 ms; interpreter-only JVM (`-Xint`): file 13.9 s → 0.5 s, in-memory 1.4 → 0.5 s; on a 32-bit ARM Android device (debuggable app): file source 23.3 → 1.2 s, in-memory 2.0 → 1.15 s warm and 4.0–4.7 → 3.0 s on the first parse in a process. + ## [0.62.2] - 2026-10-24 Headline: **the dense FP32 `x · Wᵀ` product on the JVM native lane uses every core and, on x86-64, AVX2/FMA** diff --git a/skainet-io/skainet-io-core/src/jvmTest/kotlin/sk/ainet/io/tokenizer/GgufTokenizerLoadBenchmarkTest.kt b/skainet-io/skainet-io-core/src/jvmTest/kotlin/sk/ainet/io/tokenizer/GgufTokenizerLoadBenchmarkTest.kt new file mode 100644 index 000000000..9021ae516 --- /dev/null +++ b/skainet-io/skainet-io-core/src/jvmTest/kotlin/sk/ainet/io/tokenizer/GgufTokenizerLoadBenchmarkTest.kt @@ -0,0 +1,56 @@ +package sk.ainet.io.tokenizer + +import sk.ainet.frontend.text.ArrayVocabulary +import sk.ainet.frontend.tokenizer.SentencePieceTokenizer +import sk.ainet.frontend.tokenizer.TextFrontEnds +import sk.ainet.io.ByteArrayRandomAccessSource +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.gguf.StreamingGGUFReader +import java.io.File +import kotlin.test.Test + +/** + * Timing harness (not an assertion) for loading a tokenizer from a GGUF file: metadata parse from memory and from the + * file, then the tokenizer build and its larger pieces. Set `SKAINET_GGUF_FILE` to a GGUF with tokenizer metadata; + * without it the test prints a skip notice and passes. + * + * SKAINET_GGUF_FILE=/path/model.gguf ./gradlew :skainet-io:skainet-io-core:jvmTest --tests '*GgufTokenizerLoadBenchmarkTest*' + */ +class GgufTokenizerLoadBenchmarkTest { + + private fun median(runs: Int, block: () -> Unit): Long { + repeat(2) { block() } + val ms = LongArray(runs) { val t = System.nanoTime(); block(); (System.nanoTime() - t) / 1_000_000 } + ms.sort() + return ms[runs / 2] + } + + @Test + fun timeTokenizerLoad() { + val path = System.getenv("SKAINET_GGUF_FILE") + if (path.isNullOrBlank() || !File(path).isFile) { + println("[skip] GgufTokenizerLoadBenchmarkTest: set SKAINET_GGUF_FILE to a GGUF file") + return + } + val file = File(path) + val bytes = file.readBytes() + val runs = 5 + val fields = StreamingGGUFReader.open(ByteArrayRandomAccessSource(bytes)).use { it.fields } + val tokens = (fields["tokenizer.ggml.tokens"] as List<*>).filterIsInstance() + + val readMemory = median(runs) { StreamingGGUFReader.open(ByteArrayRandomAccessSource(bytes)).use { it.fields } } + val readFile = median(3) { JvmRandomAccessSource.open(file).use { src -> StreamingGGUFReader.open(src).use { it.fields } } } + val copyTokens = median(runs) { (fields["tokenizer.ggml.tokens"] as List<*>).filterIsInstance() } + val copyScores = median(runs) { (fields["tokenizer.ggml.scores"] as List<*>).mapNotNull { (it as? Number)?.toFloat() } } + val vocabulary = median(runs) { ArrayVocabulary(tokens) } + val sentencePiece = median(runs) { SentencePieceTokenizer.fromGgufFields(fields) } + val frontEnd = median(runs) { TextFrontEnds.fromGguf(fields) } + + println( + "[bench] ${file.name} ${bytes.size} bytes, ${tokens.size} tokens, median of $runs (file: 3) in ms: " + + "read.memory=$readMemory read.file=$readFile | build.frontEnd=$frontEnd " + + "(sentencePiece.fromGgufFields=$sentencePiece, ArrayVocabulary=$vocabulary, " + + "tokens copy=$copyTokens, scores copy=$copyScores)" + ) + } +} diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMetadataCursor.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMetadataCursor.kt new file mode 100644 index 000000000..2c32cf6d5 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMetadataCursor.kt @@ -0,0 +1,177 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.RandomAccessSource + +/** + * Sequential little-endian reader over the GGUF header, KV and tensor-info sections. + * + * Reads the source in [WINDOW]-sized blocks through the buffer form of [RandomAccessSource.readAt] and decodes values + * from that block, so a primitive costs a few array reads instead of one `readAt` call (and, on a file source, one + * system call) per value. A tokenizer GGUF has arrays of ~262k strings, floats and ints; read one value at a time that + * is over a million source calls and allocations. + */ +internal class GgufMetadataCursor(private val source: RandomAccessSource, start: Long = 0L) { + + private var buf = ByteArray(WINDOW) + private var bufStart = start // file offset of buf[0] + private var bufLen = 0 // valid bytes in buf + private var p = 0 // read index in buf + + /** File offset of the next byte to read. */ + val position: Long get() = bufStart + p + + /** Make [n] bytes available at [p], refilling (and growing) the block when they are not. */ + private fun require(n: Int) { + if (bufLen - p >= n) return + val at = position + if (n > buf.size) buf = ByteArray(maxOf(n, buf.size * 2)) + var filled = 0 + val want = minOf(buf.size.toLong(), source.size - at).toInt() + while (filled < want) { + val r = source.readAt(at + filled, buf, filled, want - filled) + if (r <= 0) break + filled += r + } + bufStart = at + bufLen = filled + p = 0 + if (filled < n) throw IllegalArgumentException("Unexpected end of GGUF metadata: need $n bytes at offset $at, ${source.size - at} left") + } + + fun byte(): Byte { require(1); return buf[p++] } + + fun short(): Short { require(2); val v = (buf[p].toInt() and 0xFF) or ((buf[p + 1].toInt() and 0xFF) shl 8); p += 2; return v.toShort() } + + fun int(): Int { + require(4) + val b = buf; val i = p + p += 4 + return (b[i].toInt() and 0xFF) or ((b[i + 1].toInt() and 0xFF) shl 8) or + ((b[i + 2].toInt() and 0xFF) shl 16) or ((b[i + 3].toInt() and 0xFF) shl 24) + } + + fun long(): Long { + require(8) + val lo = int().toLong() and 0xFFFFFFFFL + val hi = int().toLong() + return lo or (hi shl 32) + } + + fun uint(): UInt = int().toUInt() + + fun ulong(): ULong = long().toULong() + + /** A GGUF string: u64 byte length, then UTF-8 bytes. */ + fun string(): String { + val length = checkedLength(ulong(), "string") + require(length) + val s = buf.decodeToString(p, p + length) + p += length + return s + } + + // Hot arrays (a tokenizer has ~262k tokens, scores and token types) are decoded straight into the output list, one + // block of at most [BLOCK] elements per call of a small method, with no per-element calls into the primitive + // readers. Short calls give a JIT runtime a method it can compile after the first blocks instead of one long loop + // that runs interpreted for its whole call (measured on a debuggable Android app: first parse ≈ 10 % faster). + + /** [count] GGUF strings appended to [out]; the same values as [string] per element. */ + fun strings(count: Int, out: MutableList) { + var left = count + while (left > 0) { + val done = stringsInBlock(minOf(left, BLOCK), out) + if (done > 0) { + left -= done + } else { // the next string straddles the block end (or is invalid): refill through string() + out.add(string()) + left-- + } + } + } + + /** Strings whose length and bytes lie wholly inside the current block, up to [max]; 0 when the next one does not. */ + private fun stringsInBlock(max: Int, out: MutableList): Int { + val b = buf + val end = bufLen + var i = p + var n = 0 + while (n < max && end - i >= 8) { + val lo = (b[i].toInt() and 0xFF) or ((b[i + 1].toInt() and 0xFF) shl 8) or + ((b[i + 2].toInt() and 0xFF) shl 16) or ((b[i + 3].toInt() and 0xFF) shl 24) + val hi = (b[i + 4].toInt() and 0xFF) or ((b[i + 5].toInt() and 0xFF) shl 8) or + ((b[i + 6].toInt() and 0xFF) shl 16) or ((b[i + 7].toInt() and 0xFF) shl 24) + if (hi != 0 || lo < 0 || end - i - 8 < lo) break // too long or not wholly in this block: string() handles it + out.add(b.decodeToString(i + 8, i + 8 + lo)) + i += 8 + lo + n++ + } + p = i + return n + } + + /** [count] 32-bit elements appended to [out] as `Int` ([INT32]), `UInt` ([UINT32]) or `Float` ([FLOAT32]). */ + fun int32s(count: Int, kind: Int, out: MutableList) { + var left = count + while (left > 0) { + require(4) + left -= int32sInBlock(minOf(left, BLOCK, (bufLen - p) / 4), kind, out) + } + } + + private fun int32sInBlock(n: Int, kind: Int, out: MutableList): Int { + val b = buf + var q = p + for (k in 0 until n) { + val v = (b[q].toInt() and 0xFF) or ((b[q + 1].toInt() and 0xFF) shl 8) or + ((b[q + 2].toInt() and 0xFF) shl 16) or ((b[q + 3].toInt() and 0xFF) shl 24) + out.add( + when (kind) { + FLOAT32 -> Float.fromBits(v) + UINT32 -> v.toUInt() + else -> v + } + ) + q += 4 + } + p = q + return n + } + + /** An array or string length as an Int index, or a clear error for a corrupt or oversized count. */ + fun checkedLength(value: ULong, what: String): Int { + if (value > Int.MAX_VALUE.toULong()) throw IllegalArgumentException("GGUF $what length $value at offset ${position - 8} is too large") + return value.toInt() + } + + /** Fail before allocating when [count] elements of at least [minBytes] each cannot fit in the rest of the source. */ + fun checkFits(count: Int, minBytes: Int, what: String) { + val left = source.size - position + if (count.toLong() * minBytes > left) throw IllegalArgumentException("GGUF $what of $count elements at offset $position needs ${count.toLong() * minBytes} bytes, $left left") + } + + // Bulk reads for the other fixed-size array elements (small arrays in practice): into a primitive array. + + fun bytes(count: Int): ByteArray { + val out = ByteArray(count) + var done = 0 + while (done < count) { + if (bufLen - p == 0) require(1) + val n = minOf(count - done, bufLen - p) + buf.copyInto(out, done, p, p + n) + p += n; done += n + } + return out + } + + fun shorts(count: Int): ShortArray = ShortArray(count).also { out -> for (i in 0 until count) out[i] = short() } + + fun longs(count: Int): LongArray = LongArray(count).also { out -> for (i in 0 until count) out[i] = long() } + + companion object { + private const val WINDOW = 64 * 1024 + private const val BLOCK = 2048 + const val INT32 = 0 + const val UINT32 = 1 + const val FLOAT32 = 2 + } +} diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt index ab23b3010..7ae0c8715 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGGUFReader.kt @@ -202,27 +202,23 @@ public class StreamingGGUFReader private constructor( // ========== Parsing Implementation ========== private fun parse() { - var offset = 0L + val cursor = GgufMetadataCursor(source) // 1. Verify magic number (4 bytes) - val magic = readUInt(offset) + val magic = cursor.uint() if (magic != GGUF_MAGIC) { throw IllegalArgumentException("Invalid GGUF magic: expected 0x${GGUF_MAGIC.toString(16)}, got 0x${magic.toString(16)}") } - offset += 4 // 2. Read version (4 bytes) - version = readUInt(offset) + version = cursor.uint() if (version.toInt() !in READER_SUPPORTED_VERSIONS) { throw IllegalArgumentException("Unsupported GGUF version: $version (supported: $READER_SUPPORTED_VERSIONS)") } - offset += 4 // 3. Read tensor count and KV count (8 + 8 bytes) - tensorCount = readULong(offset) - offset += 8 - kvCount = readULong(offset) - offset += 8 + tensorCount = cursor.ulong() + kvCount = cursor.ulong() // Store counts in fields for compatibility fields["GGUF.tensor_count"] = tensorCount @@ -230,7 +226,7 @@ public class StreamingGGUFReader private constructor( fields["GGUF.version"] = version // 4. Parse KV metadata section - offset = parseKVSection(offset, kvCount.toInt()) + parseKVSection(cursor, kvCount.toInt()) // 5. Check for alignment override val alignmentValue = fields["general.alignment"] @@ -241,9 +237,10 @@ public class StreamingGGUFReader private constructor( } // 6. Parse tensor info section - offset = parseTensorInfoSection(offset, tensorCount.toInt()) + parseTensorInfoSection(cursor, tensorCount.toInt()) // 7. Calculate data offset with alignment + var offset = cursor.position val padding = (offset % alignment).toInt() if (padding != 0) { offset += alignment - padding @@ -256,27 +253,13 @@ public class StreamingGGUFReader private constructor( } } - private fun parseKVSection(startOffset: Long, count: Int): Long { - var offset = startOffset - + private fun parseKVSection(cursor: GgufMetadataCursor, count: Int) { for (i in 0 until count) { - // Read key string - val keyLength = readULong(offset) - offset += 8 - val keyBytes = source.readAt(offset, keyLength.toInt()) - offset += keyLength.toLong() - val keyName = keyBytes.decodeToString() - - // Read value type - val valueType = readUInt(offset).toInt() - offset += 4 - + val keyName = cursor.string() + val valueType = cursor.int() val gtype = GGUFValueType.entries.find { it.value == valueType } - ?: throw IllegalArgumentException("Unknown GGUFValueType: $valueType at offset ${offset - 4}") - - // Read value - val (bytesConsumed, value) = readFieldValue(offset, gtype) - offset += bytesConsumed + ?: throw IllegalArgumentException("Unknown GGUFValueType: $valueType at offset ${cursor.position - 4}") + val value = readFieldValue(cursor, gtype) // Store in fields map if (fields.containsKey(keyName)) { @@ -285,13 +268,9 @@ public class StreamingGGUFReader private constructor( fields[keyName] = value } } - - return offset } - private fun parseTensorInfoSection(startOffset: Long, count: Int): Long { - var offset = startOffset - + private fun parseTensorInfoSection(cursor: GgufMetadataCursor, count: Int) { // First pass: collect all tensor info with relative offsets data class TensorParseInfo( val name: String, @@ -300,30 +279,19 @@ public class StreamingGGUFReader private constructor( val ggmlType: GGMLQuantizationType, val relativeOffset: Long ) - val parsedTensors = mutableListOf() + val parsedTensors = ArrayList(count) for (i in 0 until count) { - // Read tensor name - val nameLength = readULong(offset) - offset += 8 - val nameBytes = source.readAt(offset, nameLength.toInt()) - offset += nameLength.toLong() - val tensorName = nameBytes.decodeToString() - - // Read dimensions count (4 bytes) - val nDims = readUInt(offset).toInt() - offset += 4 - - // Read dimensions (8 bytes each) - val dims = mutableListOf() + val tensorName = cursor.string() + + // Dimensions: count (4 bytes), then 8 bytes each + val nDims = cursor.uint().toInt() + val dims = ArrayList(nDims) for (d in 0 until nDims) { - dims.add(readULong(offset)) - offset += 8 + dims.add(cursor.ulong()) } - // Read tensor type (4 bytes) - val tensorTypeValue = readUInt(offset).toInt() - offset += 4 + val tensorTypeValue = cursor.uint().toInt() // Use graceful fallback for unknown types val ggmlType = GGMLQuantizationType.fromValueOrUnknown(tensorTypeValue) @@ -332,9 +300,7 @@ public class StreamingGGUFReader private constructor( println("WARNING: Unknown GGML quantization type $tensorTypeValue for tensor '$tensorName'. Will treat as raw bytes.") } - // Read relative data offset (8 bytes) - val relativeDataOffset = readULong(offset) - offset += 8 + val relativeDataOffset = cursor.ulong() parsedTensors.add(TensorParseInfo( name = tensorName, @@ -385,122 +351,64 @@ public class StreamingGGUFReader private constructor( ) ) } - - return offset } - /** - * Read a field value and return (bytesConsumed, value). - */ - private fun readFieldValue(offset: Long, gtype: GGUFValueType): Pair { - return when (gtype) { - GGUFValueType.UINT8 -> 1L to source.readByteAt(offset).toUByte() - GGUFValueType.INT8 -> 1L to source.readByteAt(offset) - GGUFValueType.UINT16 -> 2L to readUShort(offset) - GGUFValueType.INT16 -> 2L to readShort(offset) - GGUFValueType.UINT32 -> 4L to readUInt(offset) - GGUFValueType.INT32 -> 4L to readInt(offset) - GGUFValueType.FLOAT32 -> 4L to readFloat(offset) - GGUFValueType.UINT64 -> 8L to readULong(offset) - GGUFValueType.INT64 -> 8L to readLong(offset) - GGUFValueType.FLOAT64 -> 8L to readDouble(offset) - GGUFValueType.BOOL -> 1L to (source.readByteAt(offset) != 0.toByte()) - GGUFValueType.STRING -> { - val length = readULong(offset) - val bytes = source.readAt(offset + 8, length.toInt()) - (8 + length.toLong()) to bytes.decodeToString() - } - GGUFValueType.ARRAY -> readArrayValue(offset) - } + /** Read one field value at the cursor. */ + private fun readFieldValue(cursor: GgufMetadataCursor, gtype: GGUFValueType): Any? = when (gtype) { + GGUFValueType.UINT8 -> cursor.byte().toUByte() + GGUFValueType.INT8 -> cursor.byte() + GGUFValueType.UINT16 -> cursor.short().toUShort() + GGUFValueType.INT16 -> cursor.short() + GGUFValueType.UINT32 -> cursor.uint() + GGUFValueType.INT32 -> cursor.int() + GGUFValueType.FLOAT32 -> Float.fromBits(cursor.int()) + GGUFValueType.UINT64 -> cursor.ulong() + GGUFValueType.INT64 -> cursor.long() + GGUFValueType.FLOAT64 -> Double.fromBits(cursor.long()) + GGUFValueType.BOOL -> cursor.byte() != 0.toByte() + GGUFValueType.STRING -> cursor.string() + GGUFValueType.ARRAY -> readArrayValue(cursor) } /** - * Read an array value. + * Read an array value into a pre-sized list with the same element types as a single value of that type. + * Elements are decoded in bulk from the cursor's block. The values are boxed here, once, into a + * plain [ArrayList]: consumers walk these lists several times (a tokenizer build reads tokens and token types three + * times), and on an interpreted runtime (a debuggable Android app) a lazily boxing view in this module cost more + * per access than the platform's compiled `ArrayList`. */ - private fun readArrayValue(startOffset: Long): Pair> { - var offset = startOffset - - // Read element type - val elementType = readUInt(offset).toInt() - offset += 4 - - // Read array length - val arrayLength = readULong(offset) - offset += 8 - + private fun readArrayValue(cursor: GgufMetadataCursor): List { + val elementType = cursor.int() + val length = cursor.checkedLength(cursor.ulong(), "array") val gtype = GGUFValueType.entries.find { it.value == elementType } ?: throw IllegalArgumentException("Unknown array element type: $elementType") - - val result = mutableListOf() - for (i in 0 until arrayLength.toInt()) { - val (bytesConsumed, value) = readFieldValue(offset, gtype) - result.add(value) - offset += bytesConsumed + // Smallest encoding per element: a string is at least its 8-byte length, a nested array its 12-byte header. + val minBytes = when (gtype) { + GGUFValueType.UINT8, GGUFValueType.INT8, GGUFValueType.BOOL -> 1 + GGUFValueType.UINT16, GGUFValueType.INT16 -> 2 + GGUFValueType.UINT32, GGUFValueType.INT32, GGUFValueType.FLOAT32 -> 4 + GGUFValueType.UINT64, GGUFValueType.INT64, GGUFValueType.FLOAT64, GGUFValueType.STRING -> 8 + GGUFValueType.ARRAY -> 12 } - - return (offset - startOffset) to result - } - - // ========== Primitive Reading Helpers ========== - - private fun readUShort(position: Long): UShort { - val bytes = source.readAt(position, 2) - return ((bytes[0].toInt() and 0xFF) or - ((bytes[1].toInt() and 0xFF) shl 8)).toUShort() - } - - private fun readShort(position: Long): Short { - val bytes = source.readAt(position, 2) - return ((bytes[0].toInt() and 0xFF) or - ((bytes[1].toInt() and 0xFF) shl 8)).toShort() - } - - private fun readUInt(position: Long): UInt { - val bytes = source.readAt(position, 4) - return ((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)).toUInt() - } - - private fun readInt(position: Long): Int { - val bytes = source.readAt(position, 4) - return (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) - } - - private fun readFloat(position: Long): Float { - return Float.fromBits(readInt(position)) - } - - private fun readULong(position: Long): ULong { - val bytes = source.readAt(position, 8) - return ((bytes[0].toLong() and 0xFF) or - ((bytes[1].toLong() and 0xFF) shl 8) or - ((bytes[2].toLong() and 0xFF) shl 16) or - ((bytes[3].toLong() and 0xFF) shl 24) or - ((bytes[4].toLong() and 0xFF) shl 32) or - ((bytes[5].toLong() and 0xFF) shl 40) or - ((bytes[6].toLong() and 0xFF) shl 48) or - ((bytes[7].toLong() and 0xFF) shl 56)).toULong() - } - - private fun readLong(position: Long): Long { - val bytes = source.readAt(position, 8) - return (bytes[0].toLong() and 0xFF) or - ((bytes[1].toLong() and 0xFF) shl 8) or - ((bytes[2].toLong() and 0xFF) shl 16) or - ((bytes[3].toLong() and 0xFF) shl 24) or - ((bytes[4].toLong() and 0xFF) shl 32) or - ((bytes[5].toLong() and 0xFF) shl 40) or - ((bytes[6].toLong() and 0xFF) shl 48) or - ((bytes[7].toLong() and 0xFF) shl 56) - } - - private fun readDouble(position: Long): Double { - return Double.fromBits(readLong(position)) + cursor.checkFits(length, minBytes, "array") + + val out = ArrayList(length) + when (gtype) { + GGUFValueType.UINT8 -> for (v in cursor.bytes(length)) out.add(v.toUByte()) + GGUFValueType.INT8 -> for (v in cursor.bytes(length)) out.add(v) + GGUFValueType.BOOL -> for (v in cursor.bytes(length)) out.add(v != 0.toByte()) + GGUFValueType.UINT16 -> for (v in cursor.shorts(length)) out.add(v.toUShort()) + GGUFValueType.INT16 -> for (v in cursor.shorts(length)) out.add(v) + GGUFValueType.UINT32 -> cursor.int32s(length, GgufMetadataCursor.UINT32, out) + GGUFValueType.INT32 -> cursor.int32s(length, GgufMetadataCursor.INT32, out) + GGUFValueType.FLOAT32 -> cursor.int32s(length, GgufMetadataCursor.FLOAT32, out) + GGUFValueType.UINT64 -> for (v in cursor.longs(length)) out.add(v.toULong()) + GGUFValueType.INT64 -> for (v in cursor.longs(length)) out.add(v) + GGUFValueType.FLOAT64 -> for (v in cursor.longs(length)) out.add(Double.fromBits(v)) + GGUFValueType.STRING -> cursor.strings(length, out) + GGUFValueType.ARRAY -> repeat(length) { out.add(readArrayValue(cursor)) } + } + return out } // ========== Closeable ========== diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StreamingGGUFReaderMetadataParityTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StreamingGGUFReaderMetadataParityTest.kt new file mode 100644 index 000000000..39a826367 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/StreamingGGUFReaderMetadataParityTest.kt @@ -0,0 +1,290 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.ByteArrayRandomAccessSource +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.gguf.export.GGUFWriter +import sk.ainet.io.gguf.export.GgufWriteRequest +import java.io.ByteArrayOutputStream +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import kotlin.test.fail + +/** + * The buffered metadata parse in [StreamingGGUFReader] must produce exactly what the earlier per-value parse produced: + * the same keys in the same order, the same value classes, the same element classes inside arrays. [LegacyFieldParser] + * is that earlier parse, kept here as the reference. + */ +class StreamingGGUFReaderMetadataParityTest { + + @Test + fun everyValueTypeMatchesThePerValueParse() { + val bytes = allTypesGguf() + assertParity(bytes) + StreamingGGUFReader.open(ByteArrayRandomAccessSource(bytes)).use { r -> + assertEquals(listOf(1, -2, Int.MAX_VALUE), r.fields["test.i32s"]) + assertEquals(listOf(1u, UInt.MAX_VALUE), r.fields["test.u32s"]) + assertEquals(listOf("", "ä€𝄞", "x"), r.fields["test.strings"]) + assertEquals(2, r.tensors.size) + assertEquals(0L, r.dataOffset % 32) + } + } + + @Test + fun tokenizerShapedArraysMatch() = assertParity(tokenizerGguf(262_172)) + + @Test + fun aSourceThatReturnsFewBytesPerReadGivesTheSameFields() { + val bytes = tokenizerGguf(20_000) + val expected = LegacyFieldParser(ByteArrayRandomAccessSource(bytes)).parse() + val actual = StreamingGGUFReader.open(TrickleSource(bytes, maxPerRead = 7)).use { it.fields } + assertSameFields(expected, actual) + } + + @Test + fun theFileSourceGivesTheSameFields() { + val file = File.createTempFile("parity_", ".gguf").apply { deleteOnExit(); writeBytes(allTypesGguf()) } + val expected = LegacyFieldParser(ByteArrayRandomAccessSource(file.readBytes())).parse() + val actual = JvmRandomAccessSource.open(file).use { src -> StreamingGGUFReader.open(src).use { it.fields } } + assertSameFields(expected, actual) + } + + @Test + fun aTruncatedFileFailsWithAClearError() { + val bytes = allTypesGguf() + val metadataEnd = StreamingGGUFReader.open(ByteArrayRandomAccessSource(bytes)).use { it.dataOffset }.toInt() + for (cut in listOf(3, 20, 40, metadataEnd / 2, metadataEnd - 40)) { + assertFailsWith("cut at $cut") { + StreamingGGUFReader.open(ByteArrayRandomAccessSource(bytes.copyOf(cut))) + } + } + } + + @Test + fun anArrayCountLargerThanTheFileFailsBeforeAllocating() { + val b = Gguf() + b.header(tensors = 0, kvs = 1) + b.key("bad.array"); b.int(GGUFValueType.ARRAY.value); b.int(GGUFValueType.FLOAT32.value); b.long(Int.MAX_VALUE.toLong()) + b.int(0) + val e = assertFailsWith { StreamingGGUFReader.open(ByteArrayRandomAccessSource(b.bytes())) } + assertTrue("left" in e.message.orEmpty(), e.message) + } + + @Test + fun parsedFieldsWriteBackThroughTheGgufWriter() { + val fields = StreamingGGUFReader.open(ByteArrayRandomAccessSource(tokenizerGguf(5_000))).use { it.fields } + val metadata = fields.filterKeys { !it.startsWith("GGUF.") }.mapValues { it.value!! } + val (_, written) = GGUFWriter.writeToByteArray(GgufWriteRequest(metadata, emptyList(), emptyMap())) + val reread = StreamingGGUFReader.open(ByteArrayRandomAccessSource(written)).use { it.fields } + for ((key, value) in metadata) assertEquals(value, reread[key], key) + } + + /** Opt-in check on a real model or tokenizer GGUF: `SKAINET_GGUF_FILE=/path/file.gguf`. */ + @Test + fun aRealGgufFileMatchesWhenConfigured() { + val path = System.getenv("SKAINET_GGUF_FILE") + if (path.isNullOrBlank() || !File(path).isFile) { + println("[skip] StreamingGGUFReaderMetadataParityTest: set SKAINET_GGUF_FILE to check a real GGUF file") + return + } + assertParity(File(path).readBytes()) + } + + // --------------------------------------------------------------------------------------------------------------- + + private fun assertParity(bytes: ByteArray) { + val expected = LegacyFieldParser(ByteArrayRandomAccessSource(bytes)).parse() + val actual = StreamingGGUFReader.open(ByteArrayRandomAccessSource(bytes)).use { it.fields } + assertSameFields(expected, actual) + } + + private fun assertSameFields(expected: Map, actual: Map) { + assertEquals(expected.keys.toList(), actual.keys.toList(), "keys and their order") + for ((key, value) in expected) assertSameValue(value, actual[key], key) + } + + private fun assertSameValue(expected: Any?, actual: Any?, path: String) { + if (expected is List<*>) { + if (actual !is List<*>) fail("$path: expected a List, got ${actual?.let { it::class }}") + assertEquals(expected.size, actual.size, "$path: size") + for (i in expected.indices) assertSameValue(expected[i], actual[i], "$path[$i]") + assertEquals(expected, actual, path) + assertEquals(expected.hashCode(), actual.hashCode(), "$path: hashCode") + return + } + assertEquals(expected?.let { it::class }, actual?.let { it::class }, "$path: class") + assertEquals(expected, actual, path) + } + + /** A source that never returns more than [maxPerRead] bytes per buffer read, like a slow stream. */ + private class TrickleSource(private val bytes: ByteArray, private val maxPerRead: Int) : RandomAccessSource { + override val size: Long get() = bytes.size.toLong() + override fun readAt(position: Long, length: Int): ByteArray = bytes.copyOfRange(position.toInt(), position.toInt() + length) + override fun readAt(position: Long, buffer: ByteArray, offset: Int, length: Int): Int { + val n = minOf(length, maxPerRead, bytes.size - position.toInt()) + if (n <= 0) return -1 + bytes.copyInto(buffer, offset, position.toInt(), position.toInt() + n) + return n + } + override fun close() {} + } + + // --------------------------------------------------------------------------------------------------------------- + // Synthetic files + + private class Gguf { + private val out = ByteArrayOutputStream() + private fun le(n: Int, fill: ByteBuffer.() -> Unit) = out.write(ByteBuffer.allocate(n).order(ByteOrder.LITTLE_ENDIAN).apply(fill).array()) + fun byte(v: Int) = out.write(v) + fun short(v: Int) = le(2) { putShort(v.toShort()) } + fun int(v: Int) = le(4) { putInt(v) } + fun long(v: Long) = le(8) { putLong(v) } + fun float(v: Float) = le(4) { putFloat(v) } + fun double(v: Double) = le(8) { putDouble(v) } + fun string(s: String) { val b = s.encodeToByteArray(); long(b.size.toLong()); out.write(b) } + fun key(k: String) = string(k) + fun header(tensors: Long, kvs: Long) { int(GGUF_MAGIC.toInt()); int(3); long(tensors); long(kvs) } + fun kv(k: String, type: GGUFValueType, value: Gguf.() -> Unit) { key(k); int(type.value); value() } + fun array(k: String, element: GGUFValueType, count: Int, each: Gguf.(Int) -> Unit) { + key(k); int(GGUFValueType.ARRAY.value); int(element.value); long(count.toLong()); for (i in 0 until count) each(i) + } + fun padTo(alignment: Int) { while (out.size() % alignment != 0) out.write(0) } + fun bytes(): ByteArray = out.toByteArray() + } + + private fun allTypesGguf(): ByteArray { + val g = Gguf() + g.header(tensors = 2, kvs = 26) + g.kv("general.architecture", GGUFValueType.STRING) { string("test") } + g.kv("test.u8", GGUFValueType.UINT8) { byte(250) } + g.kv("test.i8", GGUFValueType.INT8) { byte(0x80) } + g.kv("test.u16", GGUFValueType.UINT16) { short(0xFFFE) } + g.kv("test.i16", GGUFValueType.INT16) { short(-3) } + g.kv("test.u32", GGUFValueType.UINT32) { int(-1) } + g.kv("test.i32", GGUFValueType.INT32) { int(-7) } + g.kv("test.f32", GGUFValueType.FLOAT32) { float(-1.5f) } + g.kv("test.u64", GGUFValueType.UINT64) { long(-1L) } + g.kv("test.i64", GGUFValueType.INT64) { long(Long.MIN_VALUE) } + g.kv("test.f64", GGUFValueType.FLOAT64) { double(Math.PI) } + g.kv("test.bool", GGUFValueType.BOOL) { byte(1) } + g.kv("test.long_string", GGUFValueType.STRING) { string("ab€".repeat(40_000)) } // longer than the read block + g.array("test.u8s", GGUFValueType.UINT8, 3) { byte(listOf(0, 127, 255)[it]) } + g.array("test.i8s", GGUFValueType.INT8, 2) { byte(listOf(0x7F, 0x80)[it]) } + g.array("test.bools", GGUFValueType.BOOL, 3) { byte(listOf(0, 1, 2)[it]) } + g.array("test.u16s", GGUFValueType.UINT16, 2) { short(listOf(1, 0xFFFF)[it]) } + g.array("test.i16s", GGUFValueType.INT16, 2) { short(listOf(-1, 300)[it]) } + g.array("test.u32s", GGUFValueType.UINT32, 2) { int(listOf(1, -1)[it]) } + g.array("test.i32s", GGUFValueType.INT32, 3) { int(listOf(1, -2, Int.MAX_VALUE)[it]) } + g.array("test.f32s", GGUFValueType.FLOAT32, 4) { float(listOf(0f, -0f, Float.NaN, 1e-30f)[it]) } + g.array("test.u64s", GGUFValueType.UINT64, 2) { long(listOf(0L, -1L)[it]) } + g.array("test.i64s", GGUFValueType.INT64, 2) { long(listOf(Long.MAX_VALUE, -5L)[it]) } + g.array("test.f64s", GGUFValueType.FLOAT64, 2) { double(listOf(1.0, Double.NEGATIVE_INFINITY)[it]) } + g.array("test.strings", GGUFValueType.STRING, 3) { string(listOf("", "ä€𝄞", "x")[it]) } + g.array("test.nested", GGUFValueType.ARRAY, 3) { i -> + int(GGUFValueType.INT32.value); long(i.toLong()); repeat(i) { int(it * 10) } + } + // Two F32 tensors of 4 elements + for ((i, name) in listOf("t.a", "t.b").withIndex()) { + string(g, name); g.int(1); g.long(4); g.int(GGMLQuantizationType.F32.value); g.long(i * 32L) + } + g.padTo(32) + repeat(64) { g.byte(0) } + return g.bytes() + } + + private fun string(g: Gguf, s: String) = g.string(s) + + /** Shaped like a SentencePiece tokenizer GGUF: string tokens, float scores, int token types, a few scalars. */ + private fun tokenizerGguf(n: Int): ByteArray { + val rnd = Random(42) + val alphabet = "abcdefghijklmnopqrstuvwxyzäöüß▁<>_0123456789€" + val g = Gguf() + g.header(tensors = 0, kvs = 7) + g.kv("general.architecture", GGUFValueType.STRING) { string("gemma3") } + g.kv("tokenizer.ggml.model", GGUFValueType.STRING) { string("llama") } + g.array("tokenizer.ggml.tokens", GGUFValueType.STRING, n) { i -> + string(if (i < 3) listOf("", "", "")[i] else String(CharArray(1 + rnd.nextInt(16)) { alphabet[rnd.nextInt(alphabet.length)] })) + } + g.array("tokenizer.ggml.scores", GGUFValueType.FLOAT32, n) { g.float(-it.toFloat() / 7f) } + g.array("tokenizer.ggml.token_type", GGUFValueType.INT32, n) { int(if (it < 3) 3 else 1) } + g.kv("tokenizer.ggml.bos_token_id", GGUFValueType.UINT32) { int(2) } + g.kv("tokenizer.ggml.add_bos_token", GGUFValueType.BOOL) { byte(1) } + return g.bytes() + } + + // --------------------------------------------------------------------------------------------------------------- + + /** The per-value KV parse as it was before the buffered cursor (header and KV section only). */ + private class LegacyFieldParser(private val source: RandomAccessSource) { + private val fields = LinkedHashMap() + + fun parse(): LinkedHashMap { + var offset = 8L + val tensorCount = readULong(offset); offset += 8 + val kvCount = readULong(offset); offset += 8 + fields["GGUF.tensor_count"] = tensorCount + fields["GGUF.kv_count"] = kvCount + fields["GGUF.version"] = readUInt(4) + for (i in 0 until kvCount.toInt()) { + val keyLength = readULong(offset); offset += 8 + val keyName = source.readAt(offset, keyLength.toInt()).decodeToString(); offset += keyLength.toLong() + val valueType = readUInt(offset).toInt(); offset += 4 + val gtype = GGUFValueType.entries.first { it.value == valueType } + val (consumed, value) = readFieldValue(offset, gtype); offset += consumed + if (fields.containsKey(keyName)) fields["${keyName}_dup_$i"] = value else fields[keyName] = value + } + return fields + } + + private fun readFieldValue(offset: Long, gtype: GGUFValueType): Pair = when (gtype) { + GGUFValueType.UINT8 -> 1L to source.readByteAt(offset).toUByte() + GGUFValueType.INT8 -> 1L to source.readByteAt(offset) + GGUFValueType.UINT16 -> 2L to readInt(offset, 2).toUShort() + GGUFValueType.INT16 -> 2L to readInt(offset, 2).toShort() + GGUFValueType.UINT32 -> 4L to readUInt(offset) + GGUFValueType.INT32 -> 4L to readInt(offset, 4) + GGUFValueType.FLOAT32 -> 4L to Float.fromBits(readInt(offset, 4)) + GGUFValueType.UINT64 -> 8L to readULong(offset) + GGUFValueType.INT64 -> 8L to readULong(offset).toLong() + GGUFValueType.FLOAT64 -> 8L to Double.fromBits(readULong(offset).toLong()) + GGUFValueType.BOOL -> 1L to (source.readByteAt(offset) != 0.toByte()) + GGUFValueType.STRING -> { + val length = readULong(offset) + (8 + length.toLong()) to source.readAt(offset + 8, length.toInt()).decodeToString() + } + GGUFValueType.ARRAY -> { + var o = offset + val elementType = readUInt(o).toInt(); o += 4 + val length = readULong(o); o += 8 + val et = GGUFValueType.entries.first { it.value == elementType } + val result = mutableListOf() + for (i in 0 until length.toInt()) { + val (c, v) = readFieldValue(o, et); result.add(v); o += c + } + (o - offset) to result + } + } + + private fun readInt(position: Long, n: Int): Int { + val b = source.readAt(position, n) + var v = 0 + for (i in 0 until n) v = v or ((b[i].toInt() and 0xFF) shl (8 * i)) + return v + } + + private fun readUInt(position: Long): UInt = readInt(position, 4).toUInt() + + private fun readULong(position: Long): ULong { + val b = source.readAt(position, 8) + var v = 0L + for (i in 0 until 8) v = v or ((b[i].toLong() and 0xFF) shl (8 * i)) + return v.toULong() + } + } +}