diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e268063..e8d0e8f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Changed - **Dense FP32 SGEMM uses AVX2/FMA on x86-64** (`skainet-backend-native-cpu`): `skainet_fp32_matmul`, the kernel behind every `A · B` product the JVM hands to the native lane (a weight stored `[in, out]`, attention scores and values, training), had register-tiled bodies for wasm SIMD128 only; the x86-64 library is built for baseline x86-64, so the kernel ran a scalar/SSE2 outer-product loop without FMA that loaded and stored a row of C for every row of A. Hosts with AVX2 and FMA now select, through the cached `cpuid` probe the `x · Wᵀ` kernel introduced, a lane compiled per function for that target: a 6 × 16 tile of C in twelve eight-lane FMA accumulators over a 256-deep k-block, each sixteen-column panel of B packed once so a power-of-two `n` no longer maps every row of the panel to one L1 set, 2 × 16 and 1 × 16 tiles for the rows left over and a lane-masked stripe for the columns left over. Every output takes one fused multiply-add per `p` in ascending `p` whichever tile holds it, so sliced, threaded and whole calls stay bit-identical to each other; against the previous x86 numbers an output moves within rounding, as the NEON and wasm lanes already do. Machines without AVX2, and every other architecture, run the previous code. Measured single call on a six-core laptop, develop → this change: 63 × 1024 × 1024 21 → 281 GFLOP/s, 1024³ 48 → 311, 512³ 59 → 358, 192 × 192 × 64 (one thread) 18 → 99. +- **`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. ### Fixed 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 00000000..9021ae51 --- /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 00000000..2c32cf6d --- /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 ab23b301..7ae0c871 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 00000000..39a82636 --- /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() + } + } +}