perf(tokenizer): build a GGUF SentencePiece tokenizer in one pass over the vocab

CHANGELOG.md

@@ -5,6 +5,8 @@
 ### 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.
+- **A GGUF SentencePiece tokenizer is built in one pass over the vocabulary** (`skainet-frontend-text`, `skainet-frontend-core`): `TextFrontEnds.fromGguf` derived the token list, the token types and the special-token map three to four times per build, each a full copy of the ~262k-entry lists, and copied the scores twice. They are now taken once and passed on; the token list is used without a copy when it holds only strings, the scores go straight into a `FloatArray`, and `ArrayVocabulary` fills its map with one operation per piece (first occurrence still wins). Public API unchanged. Token ids are identical to the previous build on 26,803 texts (`GgufTokenizerIdParityTest`, env-gated). Gemma-3 tokenizer: `-Xint` JVM 1,211 → 276 ms; 32-bit ARM Android device (debuggable app): 4.5–5.4 → 3.0–3.3 s on the first build in a process, 2.6–2.8 → 1.4–1.7 s warm.
 
 ### Fixed
 

skainet-frontend/skainet-frontend-core/src/commonMain/kotlin/sk/ainet/frontend/text/Vocabulary.kt

@@ -19,8 +19,9 @@ public interface Vocabulary {
 public class ArrayVocabulary(pieces: List<String>) : Vocabulary {
     private val idToPiece: Array<String> = pieces.toTypedArray()
     private val pieceToId: Map<String, Int> = HashMap<String, Int>(pieces.size * 2).also { m ->
-        // first occurrence wins, matching how llama.cpp and `tokenizers` build their maps
-        for (i in idToPiece.indices) if (idToPiece[i] !in m) m[idToPiece[i]] = i
+        // first occurrence wins, matching how llama.cpp and `tokenizers` build their maps: inserting from the last id
+        // down, the lowest id of a duplicated piece is written last, at one map operation per piece
+        for (i in idToPiece.indices.reversed()) m[idToPiece[i]] = i
     }
 
     override val size: Int get() = idToPiece.size

skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/GgufFields.kt

@@ -48,18 +48,43 @@ internal object GgufTokenizerFields {
     const val TOKEN_TYPE_USER_DEFINED = 4
     const val TOKEN_TYPE_BYTE = 6
 
-    fun tokens(fields: Map<String, Any?>): List<String> =
-        (fields[TOKENS] as? List<*>)?.filterIsInstance<String>()
-            ?: error("$TOKENS missing or malformed")
+    // A tokenizer GGUF holds ~262k tokens and token types; each helper below makes one pass over them, and the GGUF
+    // builders call each one once per build (passing the results on) instead of re-deriving them from the fields.
 
-    fun tokenTypes(fields: Map<String, Any?>): List<Int>? =
-        (fields[TOKEN_TYPE] as? List<*>)?.mapNotNull { it.toIntFlexible() }
+    /** The token strings; the reader's list itself when every element is a `String`, without a copy. */
+    fun tokens(fields: Map<String, Any?>): List<String> {
+        val raw = fields[TOKENS] as? List<*> ?: error("$TOKENS missing or malformed")
+        @Suppress("UNCHECKED_CAST")
+        return if (raw.all { it is String }) raw as List<String> else raw.filterIsInstance<String>()
+    }
+
+    /** The token types as ints; elements that are not numbers are dropped. */
+    fun tokenTypeIds(fields: Map<String, Any?>): IntArray? {
+        val raw = fields[TOKEN_TYPE] as? List<*> ?: return null
+        val out = IntArray(raw.size)
+        var n = 0
+        for (e in raw) {
+            val v = if (e is Int) e else e.toIntFlexible() ?: continue
+            out[n++] = v
+        }
+        return if (n == out.size) out else out.copyOf(n)
+    }
+
+    /** The scores as floats, elements that are not numbers dropped; `null` when the field is missing. */
+    fun scores(fields: Map<String, Any?>): FloatArray? {
+        val raw = fields[SCORES] as? List<*> ?: return null
+        val out = FloatArray(raw.size)
+        var n = 0
+        for (e in raw) {
+            val v = if (e is Float) e else (e as? Number)?.toFloat() ?: continue
+            out[n++] = v
+        }
+        return if (n == out.size) out else out.copyOf(n)
+    }
 
-    /** CONTROL and USER_DEFINED tokens: the atomic chat-template markers, keyed by surface form. */
-    fun specialTokenMap(fields: Map<String, Any?>): Map<String, Int> {
-        val tokens = (fields[TOKENS] as? List<*>)?.filterIsInstance<String>().orEmpty()
-        val types = tokenTypes(fields).orEmpty()
-        if (tokens.isEmpty() || types.isEmpty()) return emptyMap()
+    /** CONTROL and USER_DEFINED tokens of [tokens] by [types]: the atomic chat-template markers, keyed by surface form. */
+    fun specialTokenMap(tokens: List<String>, types: IntArray?): Map<String, Int> {
+        if (tokens.isEmpty() || types == null || types.isEmpty()) return emptyMap()
         val out = LinkedHashMap<String, Int>()
         val limit = minOf(tokens.size, types.size)
         for (i in 0 until limit) {

skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/QwenByteLevelBpeTokenizer.kt

@@ -137,7 +137,7 @@ public class QwenByteLevelBpeTokenizer(
                 ?.filterIsInstance<String>()
                 ?: error("${GgufTokenizerFields.MERGES} missing — required for byte-level BPE")
             val merges = mergesRaw.map { BpeCore.parseMergeLine(it) }
-            val specials = GgufTokenizerFields.specialTokenMap(fields)
+            val specials = GgufTokenizerFields.specialTokenMap(tokens, GgufTokenizerFields.tokenTypeIds(fields))
             val policy = GgufTokenizerFields.specialTokens(fields, specials.values.toSet())
             return QwenByteLevelBpeTokenizer(
                 tokens = tokens,

skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/SentencePieceTokenizer.kt

@@ -34,29 +34,38 @@ import sk.ainet.frontend.text.Vocabulary
  * For `tokenizer.json` files with `model.type == "Unigram"` the reference algorithm is the
  * Viterbi lattice of [UnigramTokenizer]; [fromTokenizerJson] is kept for source compatibility.
  */
-public class SentencePieceTokenizer(
+public class SentencePieceTokenizer internal constructor(
     tokens: List<String>,
-    scores: List<Float>,
-    public val unknownTokenId: Int? = null,
-    bosTokenId: Int? = null,
-    eosTokenId: Int? = null,
-    public val addSpacePrefix: Boolean = true,
-    policy: SpecialTokens? = null,
-    source: String? = null,
+    private val idToScore: FloatArray,
+    public val unknownTokenId: Int?,
+    bosTokenId: Int?,
+    eosTokenId: Int?,
+    public val addSpacePrefix: Boolean,
+    policy: SpecialTokens?,
+    source: String?,
 ) : Tokenizer, RunDecoder {
 
+    public constructor(
+        tokens: List<String>,
+        scores: List<Float>,
+        unknownTokenId: Int? = null,
+        bosTokenId: Int? = null,
+        eosTokenId: Int? = null,
+        addSpacePrefix: Boolean = true,
+        policy: SpecialTokens? = null,
+        source: String? = null,
+    ) : this(tokens, scores.toFloatArray(), unknownTokenId, bosTokenId, eosTokenId, addSpacePrefix, policy, source)
+
     private val pieces: ArrayVocabulary
-    private val idToScore: FloatArray
 
     /** `byteTokenIds[b]` = vocab id of `<0xBB>`, or `-1` if absent. */
     private val byteTokenIds: IntArray
 
     init {
-        require(tokens.size == scores.size) {
-            "tokens (${tokens.size}) and scores (${scores.size}) must have the same length"
+        require(tokens.size == idToScore.size) {
+            "tokens (${tokens.size}) and scores (${idToScore.size}) must have the same length"
         }
         pieces = ArrayVocabulary(tokens)
-        idToScore = FloatArray(scores.size) { scores[it] }
         byteTokenIds = IntArray(256) { b -> pieces.tokenToId(byteTokenString(b)) ?: -1 }
     }
 
@@ -241,23 +250,32 @@ public class SentencePieceTokenizer(
          */
         public fun fromGgufFields(fields: Map<String, Any?>): SentencePieceTokenizer {
             val tokens = GgufTokenizerFields.tokens(fields)
-            val scores = (fields[GgufTokenizerFields.SCORES] as? List<*>)
-                ?.mapNotNull { (it as? Number)?.toFloat() }
+            val types = GgufTokenizerFields.tokenTypeIds(fields)
+            return fromGgufFields(fields, tokens, types, GgufTokenizerFields.specialTokenMap(tokens, types))
+        }
+
+        /** As [fromGgufFields] with the tokens, token types and special tokens already taken from [fields]. */
+        internal fun fromGgufFields(
+            fields: Map<String, Any?>,
+            tokens: List<String>,
+            types: IntArray?,
+            specials: Map<String, Int>,
+        ): SentencePieceTokenizer {
+            val scores = GgufTokenizerFields.scores(fields)
                 ?: error("${GgufTokenizerFields.SCORES} missing — required for SentencePiece")
             require(tokens.size == scores.size) { "GGUF tokens (${tokens.size}) and scores (${scores.size}) disagree" }
 
             var unknownId = fields[GgufTokenizerFields.UNK]?.toIntFlexible()
             if (unknownId == null) {
-                val idx = GgufTokenizerFields.tokenTypes(fields)?.indexOf(GgufTokenizerFields.TOKEN_TYPE_UNKNOWN) ?: -1
+                val idx = types?.indexOf(GgufTokenizerFields.TOKEN_TYPE_UNKNOWN) ?: -1
                 if (idx >= 0) unknownId = idx
             }
             val addSpacePrefix = fields[GgufTokenizerFields.ADD_SPACE_PREFIX].toBooleanFlexible(default = true)
-            val specials = GgufTokenizerFields.specialTokenMap(fields)
             val policy = GgufTokenizerFields.specialTokens(fields, specials.values.toSet()).let { if (it.unk == null) it.copy(unk = unknownId) else it }
 
             return SentencePieceTokenizer(
                 tokens = tokens,
-                scores = scores,
+                idToScore = scores,
                 unknownTokenId = unknownId,
                 bosTokenId = policy.bos,
                 eosTokenId = policy.eos,

skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/TextFrontEnds.kt

@@ -106,9 +106,11 @@ public object TextFrontEnds : FrontEndProvider {
         return when (model) {
             "gpt2", "bpe" -> QwenByteLevelBpeTokenizer.fromGgufFields(fields)
             "llama", "sentencepiece", "gemma4" -> {
-                val base = SentencePieceTokenizer.fromGgufFields(fields)
-                val specials = GgufTokenizerFields.specialTokenMap(fields)
-                if (specials.isEmpty()) base else SpecialTokenSplitter(base, specials, specialIds = specials.values.toSet())
+                val tokens = GgufTokenizerFields.tokens(fields)
+                val types = GgufTokenizerFields.tokenTypeIds(fields)
+                val specials = GgufTokenizerFields.specialTokenMap(tokens, types)
+                val base = SentencePieceTokenizer.fromGgufFields(fields, tokens, types, specials)
+                if (specials.isEmpty()) base else SpecialTokenSplitter(base, specials, specialIds = base.specialTokens.declared)
             }
             "bert", "wordpiece" -> {
                 val tokens = GgufTokenizerFields.tokens(fields)

skainet-io/skainet-io-core/src/jvmTest/kotlin/sk/ainet/io/tokenizer/GgufTokenizerIdParityTest.kt

@@ -0,0 +1,63 @@
+package sk.ainet.io.tokenizer
+
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.jsonPrimitive
+import sk.ainet.frontend.tokenizer.TextFrontEnds
+import sk.ainet.io.ByteArrayRandomAccessSource
+import sk.ainet.io.gguf.StreamingGGUFReader
+import java.io.File
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+/**
+ * Token-id parity of a GGUF tokenizer against ids recorded earlier: guards changes to the GGUF reader and the tokenizer
+ * build, which must not move a single id. Needs three files, none of them in the repo:
+ *
+ * - `SKAINET_GGUF_FILE`: a GGUF with tokenizer metadata;
+ * - `SKAINET_TOKENIZER_CORPUS`: JSON lines, one JSON string per line, the texts to encode;
+ * - `SKAINET_TOKENIZER_IDS`: the recorded ids. When the file does not exist the test writes it (record a baseline with
+ *   the build to compare against), otherwise it compares line by line.
+ *
+ * Besides `encode(text)` per text it records the vocabulary round trip over every id and the special-token policy.
+ * Without the first two variables the test prints a skip notice and passes.
+ */
+class GgufTokenizerIdParityTest {
+
+    @Test
+    fun idsMatchRecordedBaseline() {
+        val gguf = System.getenv("SKAINET_GGUF_FILE")?.let(::File)
+        val corpus = System.getenv("SKAINET_TOKENIZER_CORPUS")?.let(::File)
+        val idsPath = System.getenv("SKAINET_TOKENIZER_IDS")
+        if (gguf?.isFile != true || corpus?.isFile != true || idsPath.isNullOrBlank()) {
+            println("[skip] GgufTokenizerIdParityTest: set SKAINET_GGUF_FILE, SKAINET_TOKENIZER_CORPUS and SKAINET_TOKENIZER_IDS")
+            return
+        }
+        val fields = StreamingGGUFReader.open(ByteArrayRandomAccessSource(gguf.readBytes())).use { it.fields }
+        val tokenizer = TextFrontEnds.fromGguf(fields)
+        val vocab = tokenizer.vocabulary
+
+        val lines = ArrayList<String>()
+        lines += "specialTokens ${tokenizer.specialTokens}"
+        lines += "vocab.size ${vocab.size}"
+        var roundTrip = 0L
+        for (id in 0 until vocab.size) {
+            val piece = vocab.idToToken(id)
+            roundTrip = roundTrip * 31 + (piece?.let { vocab.tokenToId(it) } ?: -2) + 7 * (piece?.hashCode() ?: 0)
+        }
+        lines += "vocab.roundTrip $roundTrip"
+        corpus.forEachLine { line ->
+            if (line.isNotBlank()) lines += tokenizer.encode(Json.parseToJsonElement(line).jsonPrimitive.content).joinToString(" ")
+        }
+
+        val ids = File(idsPath)
+        if (!ids.exists()) {
+            ids.writeText(lines.joinToString("\n", postfix = "\n"))
+            println("[record] GgufTokenizerIdParityTest: wrote ${lines.size} lines to $ids")
+            return
+        }
+        val expected = ids.readLines()
+        assertEquals(expected.size, lines.size, "line count")
+        for (i in expected.indices) assertEquals(expected[i], lines[i], "line ${i + 1}")
+        println("[parity] GgufTokenizerIdParityTest: ${lines.size} lines identical (${lines.size - 3} texts)")
+    }
+}

skainet-io/skainet-io-core/src/jvmTest/kotlin/sk/ainet/io/tokenizer/GgufTokenizerLoadBenchmarkTest.kt

@@ -0,0 +1,53 @@
+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<String>()
+
+        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 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)"
+        )
+    }
+}