From 86d32c38500f220de45483ff047c8ead5bdb078e Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 2 Sep 2026 14:56:43 +0200 Subject: [PATCH 1/2] refactor(llm-core): SafeTensors loading collapses onto the engine loader (SKaiNET#1246) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DecoderSafeTensorsLoader.loadToMap rides SafeTensorsParametersLoader.withPolicy for every all-float checkpoint: the engine owns reading and the BF16/F16 widen-or-keep-native decision for the DTypePolicy; the family keeps only what the engine does not know — HF -> GGUF renaming, [1, dim] norm normalization, tied embeddings, and the input-major relayout of keep-native matmul weights (engine #888). The public surface is unchanged (loadToMap stays a plain function for the Java consumers: the engine load is suspend by contract but never suspends, so it is driven to completion without a coroutine runtime). The legacy Q4 + .qb export, non-float tensors and non-FP32 witnesses cannot ride the engine's single-file loader (no tensorFilter, no Q4 dtype); those files keep the pre-#1246 reader verbatim, scoped as loadLegacy. Co-Authored-By: Claude Fable 5.1 --- .../dsl/decoder/DecoderSafeTensorsLoader.kt | 182 +++++++++++----- .../DecoderSafeTensorsLoaderFixtureTest.kt | 203 ++++++++++++++++++ 2 files changed, 330 insertions(+), 55 deletions(-) create mode 100644 llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoaderFixtureTest.kt diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoader.kt b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoader.kt index 9b6b781a..b09389a1 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoader.kt +++ b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoader.kt @@ -2,15 +2,14 @@ package sk.ainet.lang.nn.dsl.decoder import sk.ainet.context.ExecutionContext import sk.ainet.io.RandomAccessSource -import sk.ainet.lang.nn.dsl.decoder.GgufDecoderMetadata -import sk.ainet.lang.nn.dsl.decoder.DecoderGgufWeights -import sk.ainet.lang.nn.dsl.decoder.DecoderTensorNames import sk.ainet.io.model.DataType -import sk.ainet.io.safetensors.StreamingSafeTensorsReader +import sk.ainet.io.safetensors.SafeTensorsParametersLoader import sk.ainet.io.safetensors.StreamingSafeTensorInfo +import sk.ainet.io.safetensors.StreamingSafeTensorsReader import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.apps.llm.DTypePolicyValidation +import sk.ainet.lang.tensor.data.FloatArrayTensorData import sk.ainet.lang.tensor.data.NarrowFloatDenseTensorData import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData import sk.ainet.lang.tensor.data.NarrowFloatTensorData @@ -23,31 +22,39 @@ import sk.ainet.lang.types.DType import sk.ainet.lang.types.DTypePolicy import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 +import kotlin.coroutines.Continuation +import kotlin.coroutines.EmptyCoroutineContext +import kotlin.coroutines.startCoroutine import kotlin.math.pow import kotlin.reflect.KClass /** - * Loads LLaMA weights from HuggingFace SafeTensors format and maps them to - * the canonical GGUF tensor naming used by [LlamaWeightMapper]. + * Loads LLaMA-family weights from a HuggingFace SafeTensors file and maps them to + * the canonical GGUF tensor naming used by `LlamaWeightMapper`. * - * Handles: - * - HuggingFace → GGUF tensor name mapping - * - Q4 + .qb companion tensor dequantization to FP32 - * - BF16/F16 dequantization to FP32 (default) - * - **Narrow-float KEEP_NATIVE** when [dtypePolicy] admits BF16 or F16 - * (SKaiNET 0.38.0): keeps the on-disk 2-bytes-per-element buffer so the narrow-float - * matmul kernel can run without a 2× memory blow-up. Matmul weights are relaid - * input-major so the per-forward transpose is free — see [narrowData]. + * Reading and materialization ride the engine (SKaiNET#1246): the file goes through + * [SafeTensorsParametersLoader.withPolicy], which owns the BF16/F16 widening and + * keep-native decisions for [dtypePolicy]. What stays family-side is the policy the + * engine does not know about: + * + * - HuggingFace → GGUF tensor name mapping ([HfTensorNameMapper]); unmapped tensors are dropped * - Shape normalization ([1, dim] norms → [dim]) + * - **Input-major relayout of keep-native matmul weights** — see [narrowData] * - Tied word embeddings (output.weight = token_embd.weight) * + * One file shape cannot go through the engine yet: the legacy `Q4` + `.qb` companion export + * (a SKaiNET-specific format the engine has no dtype for), and files that carry non-float or + * F64 tensors, which the engine's single-file loader would either reject or widen instead of + * skipping. Those files take [loadLegacy], the pre-#1246 reader kept verbatim for them. + * Engine follow-up: a `tensorFilter` on the single-file loader (the sharded one has it) would + * let the standard tensors of such files ride the engine as well. + * * @param dtypePolicy declarative dtype constraint. Default [DTypePolicy.Any] * = widen everything to FP32. `Require(X)` / `Prefer(X)` / `OneOf` containing * X keeps X-encoded source tensors packed, for X in {BF16, FP16}. The two * formats are resolved independently — `Require(BF16)` still widens F16 * sources, since neither narrow format can be re-encoded as the other without - * a lossy round-trip. Mirrors the engine-side - * `SafeTensorsParametersLoader.mapPolicyToNarrow` semantics. + * a lossy round-trip. This is the engine's own `mapPolicyToNarrow` semantics. */ public class DecoderSafeTensorsLoader( private val ctx: ExecutionContext, @@ -58,10 +65,8 @@ public class DecoderSafeTensorsLoader( ) { /** - * Returns `true` iff [dtypePolicy] wants BF16 weights kept in their - * packed 2-bytes-per-element form rather than widened to FP32. - * Matches the engine-side `SafeTensorsParametersLoader.mapPolicyToNarrow` - * cases that resolve to `NarrowFloatLoadPolicy.KEEP_NATIVE`. + * Returns `true` iff [dtypePolicy] wants BF16 weights kept in their packed form — + * used by the legacy path only; the engine resolves this itself on the main path. */ private val keepBf16Native: Boolean = DTypePolicyValidation.keepsNative(dtypePolicy, BF16) @@ -73,6 +78,88 @@ public class DecoderSafeTensorsLoader( * Useful for feeding into [WeightMapper] with a [WeightNameResolver]. */ public fun loadToMap(randomAccessProvider: () -> RandomAccessSource): DecoderGgufWeights { + // Header-only pass: decide whether the engine loader can consume this file. + val header = StreamingSafeTensorsReader.open(randomAccessProvider()).use { it.tensors.toList() } + val tensors = if (engineEligible(header)) { + loadViaEngine(randomAccessProvider) + } else { + loadLegacy(randomAccessProvider) + } + + // Handle tied embeddings: reuse token_embd as output.weight + if (tiedEmbeddings && !tensors.containsKey(DecoderTensorNames.OUTPUT_WEIGHT)) { + val embedding = tensors[DecoderTensorNames.TOKEN_EMBEDDINGS] + ?: error("tie_word_embeddings=true but token embedding not found") + tensors[DecoderTensorNames.OUTPUT_WEIGHT] = embedding + println(" Tied: ${DecoderTensorNames.OUTPUT_WEIGHT} → ${DecoderTensorNames.TOKEN_EMBEDDINGS}") + } + + return DecoderGgufWeights( + metadata = metadata, + tensors = tensors + ) + } + + // ========== Engine path ========== + + /** + * The engine's single-file loader delivers every tensor in the file and requires an FP32 + * dtype witness for float sources; it has no filter, so a file is engine-eligible only when + * every tensor is one the engine materializes the way this loader always did (F32/F16/BF16). + */ + private fun engineEligible(header: List): Boolean = + dtype == FP32::class && header.all { it.dataType in ENGINE_FLOAT_TYPES } + + private fun loadViaEngine(randomAccessProvider: () -> RandomAccessSource): MutableMap> { + val tensors = linkedMapOf>() + val loader = SafeTensorsParametersLoader.withPolicy(randomAccessProvider, dtypePolicy) + runNonSuspending { + loader.load(ctx, dtype) { hfName, tensor -> + val canonicalName = HfTensorNameMapper.toCanonical(hfName) ?: return@load + val adapted = adapt(canonicalName, tensor) + tensors[canonicalName] = adapted + println(" Loaded: $hfName ${tensor.shape} → $canonicalName ${adapted.shape}") + } + } + return tensors + } + + /** + * Apply the family-side shape/layout policy to an engine-materialized tensor: norm shape + * normalization, and the input-major relayout of keep-native matmul weights. Tensors that + * need neither are returned as delivered (no copy). + */ + @Suppress("UNCHECKED_CAST") + private fun adapt(canonicalName: String, tensor: Tensor): Tensor { + val targetShape = normalizeNormShape(tensor.shape) + return when (val data = tensor.data) { + is NarrowFloatDenseTensorData -> { + val relaid = narrowData(canonicalName, targetShape, data.packedData, data.codec, existing = data) + if (relaid === data) tensor else ctx.fromData(relaid as TensorData, dtype) + } + is FloatArrayTensorData<*> -> { + if (targetShape == tensor.shape) tensor + else ctx.wrapFloatArray(targetShape, dtype, data.buffer) as Tensor + } + else -> tensor + } + } + + /** + * The engine's single-file `load` is `suspend` by interface contract but never actually + * suspends (the reader is synchronous), so it can be driven to completion without a + * coroutine runtime — which keeps [loadToMap] a plain function for the Java consumers + * (`KLlamaJava`) and the non-suspending network loaders that call it. + */ + private fun runNonSuspending(block: suspend () -> R): R { + var outcome: Result? = null + block.startCoroutine(Continuation(EmptyCoroutineContext) { outcome = it }) + return (outcome ?: error("engine SafeTensors load suspended unexpectedly")).getOrThrow() + } + + // ========== Legacy path (Q4 + .qb export, non-float tensors, non-FP32 witnesses) ========== + + private fun loadLegacy(randomAccessProvider: () -> RandomAccessSource): MutableMap> { val tensors = mutableMapOf>() StreamingSafeTensorsReader.open(randomAccessProvider()).use { reader -> @@ -101,17 +188,6 @@ public class DecoderSafeTensorsLoader( val bytes = reader.loadTensorData(info) val targetShape = normalizeNormShape(info.shape) if (keepBf16Native) { - // KEEP_NATIVE: wrap the packed 2-bytes-per-element - // BF16 buffer as `Bf16DenseTensorData`. The matmul - // dispatch in `DefaultCpuOpsJvm` (SKaiNET 0.38.0) - // detects `NarrowFloatTensorData` at runtime and - // picks the kernel by codec — avoiding the 2× memory - // inflation of the FP32 dequant path. - // - // The declared dtype generic stays `T` (typically - // FP32) because consumers don't care about the - // physical encoding — the get/set surface still - // returns Float. val data = narrowData(canonicalName, targetShape, bytes, Bf16Codec) @Suppress("UNCHECKED_CAST") ctx.fromData(data as TensorData, dtype) as Tensor @@ -125,11 +201,6 @@ public class DecoderSafeTensorsLoader( val bytes = reader.loadTensorData(info) val targetShape = normalizeNormShape(info.shape) if (keepFp16Native) { - // Mirrors the BF16 arm above. Distinct from it on purpose: - // both formats are 2 bytes per element, so handing F16 bytes - // to the BF16 decode would not throw — it would produce - // plausible-looking wrong numbers. The codec carried by - // `Fp16DenseTensorData` is what keeps the dispatch honest. val data = narrowData(canonicalName, targetShape, bytes, Fp16Codec) @Suppress("UNCHECKED_CAST") ctx.fromData(data as TensorData, dtype) as Tensor @@ -156,22 +227,9 @@ public class DecoderSafeTensorsLoader( println(" Loaded: ${info.name} (${info.dtype} ${info.shape}) → $canonicalName ${tensor.shape}") } } - - // Handle tied embeddings: reuse token_embd as output.weight - if (tiedEmbeddings && !tensors.containsKey(DecoderTensorNames.OUTPUT_WEIGHT)) { - val embedding = tensors[DecoderTensorNames.TOKEN_EMBEDDINGS] - ?: error("tie_word_embeddings=true but token embedding not found") - tensors[DecoderTensorNames.OUTPUT_WEIGHT] = embedding - println(" Tied: ${DecoderTensorNames.OUTPUT_WEIGHT} → ${DecoderTensorNames.TOKEN_EMBEDDINGS}") - } - - return DecoderGgufWeights( - metadata = metadata, - tensors = tensors - ) + return tensors } - /** * Build the KEEP_NATIVE storage for one narrow-float tensor, choosing its byte layout. * @@ -191,18 +249,22 @@ public class DecoderSafeTensorsLoader( * Note this also covers tied embeddings: `output.weight` aliases `token_embd`, so in the * tied case the output projection stays row-major too and forgoes the transpose win. That * is deliberate — one shared buffer cannot be optimal for both access patterns. + * + * @param existing the engine-delivered row-major data, returned as-is when it already has + * the right layout and shape (so the engine path adds no copy for norms/embeddings). */ private fun narrowData( canonicalName: String, shape: Shape, bytes: ByteArray, codec: NarrowFloatCodec, + existing: NarrowFloatDenseTensorData? = null, ): NarrowFloatTensorData { val isGatheredEmbedding = canonicalName == DecoderTensorNames.TOKEN_EMBEDDINGS - return if (shape.rank == 2 && !isGatheredEmbedding) { - NarrowFloatInputMajorTensorData.fromRowMajor(shape, bytes, codec) - } else { - NarrowFloatDenseTensorData(shape, bytes, codec) + return when { + shape.rank == 2 && !isGatheredEmbedding -> NarrowFloatInputMajorTensorData.fromRowMajor(shape, bytes, codec) + existing != null && existing.shape == shape -> existing + else -> NarrowFloatDenseTensorData(shape, bytes, codec) } } @@ -226,6 +288,11 @@ public class DecoderSafeTensorsLoader( } } + private fun normalizeNormShape(shape: Shape): Shape { + val dims = shape.dimensions + return if (dims.size == 2 && dims[0] == 1) Shape(dims[1]) else shape + } + // ========== Q4 Dequantization ========== /** @@ -290,7 +357,7 @@ public class DecoderSafeTensorsLoader( return out } - // ========== BF16/F16 Dequantization ========== + // ========== BF16/F16 Dequantization (legacy path) ========== private fun dequantBF16(bytes: ByteArray): FloatArray { val out = FloatArray(bytes.size / 2) @@ -347,6 +414,11 @@ public class DecoderSafeTensorsLoader( } return out } + + private companion object { + /** Source dtypes the engine's single-file loader materializes exactly as this loader always did. */ + val ENGINE_FLOAT_TYPES: Set = setOf(DataType.FLOAT32, DataType.FLOAT16, DataType.BFLOAT16) + } } /** diff --git a/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoaderFixtureTest.kt b/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoaderFixtureTest.kt new file mode 100644 index 00000000..bbcf702c --- /dev/null +++ b/llm-core/src/jvmTest/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderSafeTensorsLoaderFixtureTest.kt @@ -0,0 +1,203 @@ +package sk.ainet.lang.nn.dsl.decoder + +import kotlinx.io.Buffer +import kotlinx.io.readByteArray +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.safetensors.SafeTensorsWriter +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNotNull +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * End-to-end fixture for the engine-backed [DecoderSafeTensorsLoader] (SKaiNET#1246): a + * synthetic single-file LLaMA checkpoint written with the engine's `SafeTensorsWriter`, run + * through the engine `SafeTensorsParametersLoader` and this loader's family-side policy — + * HF → GGUF renaming, `[1, dim]` norm normalization, tied embeddings, keep-native relayout. + * The legacy `Q4` + `.qb` export is pinned separately because it cannot ride the engine. + */ +class DecoderSafeTensorsLoaderFixtureTest { + + private val dim = 4 + private val ffn = 4 + private val vocab = 4 + + private val metadata = GgufDecoderMetadata( + architecture = "llama", + embeddingLength = dim, + contextLength = 8, + blockCount = 1, + headCount = 1, + kvHeadCount = 1, + feedForwardLength = ffn, + ropeDimensionCount = dim, + vocabSize = vocab, + ) + + @Test + fun `engine path maps every slot, normalizes norms, ties the output and drops unmapped tensors`() { + val file = writeHfCheckpoint() + val t = load(file, DTypePolicy.Any, tied = true) + + val expected = listOf( + DecoderTensorNames.TOKEN_EMBEDDINGS, DecoderTensorNames.OUTPUT_NORM, DecoderTensorNames.OUTPUT_WEIGHT, + DecoderTensorNames.attnNorm(0), DecoderTensorNames.attnQ(0), DecoderTensorNames.attnK(0), + DecoderTensorNames.attnV(0), DecoderTensorNames.attnOut(0), DecoderTensorNames.ffnNorm(0), + DecoderTensorNames.ffnGate(0), DecoderTensorNames.ffnUp(0), DecoderTensorNames.ffnDown(0), + ) + for (slot in expected) assertNotNull(t[slot], "slot $slot must be populated; got ${t.keys}") + assertEquals(expected.toSet(), t.keys, "no extra slots — the rotary decoy must not be delivered") + assertFalse(t.keys.any { "rotary" in it }) + + // [1, dim] norms are normalized to [dim]; matmul weights keep [out, in]. + assertEquals(Shape(dim), t[DecoderTensorNames.OUTPUT_NORM]!!.shape) + assertEquals(Shape(dim), t[DecoderTensorNames.attnNorm(0)]!!.shape) + assertEquals(Shape(dim, dim), t[DecoderTensorNames.attnQ(0)]!!.shape) + assertEquals(Shape(vocab, dim), t[DecoderTensorNames.TOKEN_EMBEDDINGS]!!.shape) + + // Tied: output.weight is the very same tensor as token_embd. + assertSame(t[DecoderTensorNames.TOKEN_EMBEDDINGS], t[DecoderTensorNames.OUTPUT_WEIGHT]) + + // Values round-trip exactly (bf16-representable fixtures) and Any widens to FP32. + assertContentEquals(values(dim * dim, Q_BASE), floats(t[DecoderTensorNames.attnQ(0)]!!)) + assertContentEquals(values(dim, NORM_BASE), floats(t[DecoderTensorNames.OUTPUT_NORM]!!)) + assertContentEquals(values(vocab * dim, EMBED_BASE), floats(t[DecoderTensorNames.TOKEN_EMBEDDINGS]!!)) + assertTrue(t[DecoderTensorNames.attnQ(0)]!!.data is FloatArrayTensorData<*>, "Any must widen BF16") + } + + @Test + fun `Require BF16 keeps matmul weights packed input-major and leaves F32 norms dense`() { + val file = writeHfCheckpoint() + val t = load(file, DTypePolicy.Require(BF16), tied = false) + + val q = t[DecoderTensorNames.attnQ(0)]!!.data + assertTrue(q is NarrowFloatInputMajorTensorData, "a matmul weight must be relaid input-major, got ${q::class.simpleName}") + assertEquals(Bf16Codec, (q as NarrowFloatTensorData).codec) + assertContentEquals(values(dim * dim, Q_BASE), t[DecoderTensorNames.attnQ(0)]!!.data.copyToFloatArray()) + + // The gathered embedding stays row-major (still packed), norms are dense FP32. + val embed = t[DecoderTensorNames.TOKEN_EMBEDDINGS]!!.data + assertTrue(embed is NarrowFloatTensorData && embed !is NarrowFloatInputMajorTensorData) + assertTrue(t[DecoderTensorNames.OUTPUT_NORM]!!.data is FloatArrayTensorData<*>) + assertEquals(Shape(dim), t[DecoderTensorNames.OUTPUT_NORM]!!.shape) + } + + @Test + fun `legacy Q4 plus qb export still loads through the retained reader`() { + // 4x4 q_proj: nibbles 1..4 per row, one FP32 scale per row → q * scale. + val rows = dim; val cols = dim + val q4 = ByteArray(rows * cols / 2) + for (r in 0 until rows) for (c in 0 until cols) { + val flat = r * cols + c + val nib = (c + 1) and 0x0F + val idx = flat / 2 + q4[idx] = if (flat % 2 == 0) (q4[idx].toInt() or nib).toByte() else (q4[idx].toInt() or (nib shl 4)).toByte() + } + val scales = floatArrayOf(0.5f, 0.25f, 1.0f, 2.0f) + val file = writeRaw( + listOf( + Entry("model.layers.0.self_attn.q_proj.weight", "Q4", listOf(rows, cols), q4), + Entry("model.layers.0.self_attn.q_proj.weight.qb", "F32", listOf(rows, 1), f32Bytes(scales)), + ) + ) + val t = load(file, DTypePolicy.Any, tied = false) + val expected = FloatArray(rows * cols) { (it % cols + 1).toFloat() * scales[it / cols] } + assertContentEquals(expected, floats(t[DecoderTensorNames.attnQ(0)]!!)) + assertFalse(t.containsKey("${DecoderTensorNames.attnQ(0)}.qb")) + } + + // ---- fixtures ---- + + private fun load(file: File, policy: DTypePolicy, tied: Boolean): Map> { + val loader = DecoderSafeTensorsLoader( + ctx = DirectCpuExecutionContext(), + dtype = FP32::class, + metadata = metadata, + tiedEmbeddings = tied, + dtypePolicy = policy, + ) + val provider: () -> RandomAccessSource = { JvmRandomAccessSource.open(file) } + return loader.loadToMap(provider).tensors + } + + /** A 1-layer LLaMA-style HF checkpoint: BF16 projections + embedding, F32 norms stored `[1, dim]`. */ + private fun writeHfCheckpoint(): File { + val buffer = Buffer() + SafeTensorsWriter.write(buffer) { + tensorBF16("model.embed_tokens.weight", listOf(vocab.toLong(), dim.toLong()), values(vocab * dim, EMBED_BASE)) + tensorF32("model.norm.weight", listOf(1L, dim.toLong()), values(dim, NORM_BASE)) + val l = "model.layers.0" + tensorF32("$l.input_layernorm.weight", listOf(1L, dim.toLong()), values(dim, LN_BASE)) + tensorF32("$l.post_attention_layernorm.weight", listOf(1L, dim.toLong()), values(dim, 2.5f)) + tensorBF16("$l.self_attn.q_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, Q_BASE)) + tensorBF16("$l.self_attn.k_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, 0.125f)) + tensorBF16("$l.self_attn.v_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, 0.25f)) + tensorBF16("$l.self_attn.o_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, 0.375f)) + tensorBF16("$l.mlp.gate_proj.weight", listOf(ffn.toLong(), dim.toLong()), values(ffn * dim, 0.5f)) + tensorBF16("$l.mlp.up_proj.weight", listOf(ffn.toLong(), dim.toLong()), values(ffn * dim, 0.625f)) + tensorBF16("$l.mlp.down_proj.weight", listOf(dim.toLong(), ffn.toLong()), values(dim * ffn, 0.75f)) + // Unmapped float tensor: materialized by the engine, dropped by the family mapping. + tensorF32("model.rotary_emb.inv_freq", listOf(2L), floatArrayOf(1.0f, 0.5f)) + } + val file = Files.createTempFile("decoder_st_fixture", ".safetensors").toFile().also { it.deleteOnExit() } + file.writeBytes(buffer.readByteArray()) + return file + } + + private data class Entry(val name: String, val dtype: String, val shape: List, val bytes: ByteArray) + + /** Hand-written header for dtypes the writer has no typed helper for (the legacy Q4 export). */ + private fun writeRaw(entries: List): File { + val header = StringBuilder("{") + var offset = 0L + entries.forEachIndexed { i, e -> + if (i > 0) header.append(",") + header.append("\"${e.name}\":{\"dtype\":\"${e.dtype}\",\"shape\":[${e.shape.joinToString(",")}],\"data_offsets\":[$offset,${offset + e.bytes.size}]}") + offset += e.bytes.size + } + header.append("}") + val headerBytes = header.toString().toByteArray(Charsets.UTF_8) + val file = Files.createTempFile("decoder_st_q4", ".safetensors").toFile().also { it.deleteOnExit() } + file.outputStream().use { out -> + out.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(headerBytes.size.toLong()).array()) + out.write(headerBytes) + entries.forEach { out.write(it.bytes) } + } + return file + } + + private fun f32Bytes(values: FloatArray): ByteArray { + val bb = ByteBuffer.allocate(values.size * 4).order(ByteOrder.LITTLE_ENDIAN) + values.forEach { bb.putFloat(it) } + return bb.array() + } + + /** bf16-exact values: base + i/8 stays within 7 mantissa bits for these tiny extents. */ + private fun values(n: Int, base: Float) = FloatArray(n) { base + (it % 16) * 0.125f } + private fun floats(t: Tensor): FloatArray = t.data.copyToFloatArray() + + private companion object { + const val EMBED_BASE = 1.0f + const val NORM_BASE = 2.0f + const val Q_BASE = -1.0f + const val LN_BASE = 0.5f + } +} From 3701de6c9c43c83d9bab97ebfeacd593b29d6828 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 2 Sep 2026 14:56:43 +0200 Subject: [PATCH 2/2] test(llama): synthetic HF checkpoint through the engine-backed SafeTensors load into LlamaRuntimeWeights (SKaiNET#1246) DecoderSafeTensorsLoaderLlama is a 10-line mapping extension with no materialization of its own; this pins the engine loader -> family renaming -> LlamaWeightMapper contract end to end for both policy arms. Co-Authored-By: Claude Fable 5.1 --- ...ecoderSafeTensorsLoaderLlamaFixtureTest.kt | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderLlamaFixtureTest.kt diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderLlamaFixtureTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderLlamaFixtureTest.kt new file mode 100644 index 00000000..d103ae6e --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/DecoderSafeTensorsLoaderLlamaFixtureTest.kt @@ -0,0 +1,112 @@ +package sk.ainet.models.llama + +import kotlinx.io.Buffer +import kotlinx.io.readByteArray +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.safetensors.SafeTensorsWriter +import sk.ainet.lang.nn.dsl.decoder.DecoderSafeTensorsLoader +import sk.ainet.lang.nn.dsl.decoder.GgufDecoderMetadata +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.file.Files +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * The llama half of the engine-backed SafeTensors load (SKaiNET#1246): a synthetic HF + * checkpoint through [DecoderSafeTensorsLoader.load] into structured [LlamaRuntimeWeights], + * i.e. the engine loader → family renaming → [LlamaWeightMapper] shape contract, end to end. + */ +class DecoderSafeTensorsLoaderLlamaFixtureTest { + + private val dim = 4 + private val ffn = 8 + private val vocab = 6 + + private val metadata = GgufDecoderMetadata( + architecture = "llama", + embeddingLength = dim, + contextLength = 8, + blockCount = 1, + headCount = 1, + kvHeadCount = 1, + feedForwardLength = ffn, + ropeDimensionCount = dim, + vocabSize = vocab, + ) + + @Test + fun `runtime weights come out of the engine path with the mapper's shapes and values`() { + val file = writeHfCheckpoint() + val weights = loader(DTypePolicy.Any).load { JvmRandomAccessSource.open(file) } + + assertEquals(1, weights.layers.size) + assertEquals(Shape(vocab, dim), weights.tokenEmbedding.shape) + assertEquals(Shape(dim), weights.outputNorm.shape, "[1, dim] norm must be normalized before mapping") + assertSame(weights.tokenEmbedding, weights.outputWeight, "tied embeddings") + val layer = weights.layers[0] + assertEquals(Shape(dim, dim), layer.wq.shape) + assertEquals(Shape(ffn, dim), layer.ffnGate.shape) + assertEquals(Shape(dim, ffn), layer.ffnDown.shape) + assertContentEquals(values(dim * dim, Q_BASE), layer.wq.data.copyToFloatArray()) + assertContentEquals(values(dim, LN_BASE), layer.attnNorm.data.copyToFloatArray()) + assertTrue(layer.wq.data is FloatArrayTensorData<*>, "Any widens BF16 to FP32") + } + + @Test + fun `Require BF16 reaches the runtime weights as input-major packed projections`() { + val file = writeHfCheckpoint() + val weights = loader(DTypePolicy.Require(BF16)).load { JvmRandomAccessSource.open(file) } + val layer = weights.layers[0] + assertTrue(layer.wq.data is NarrowFloatInputMajorTensorData, "got ${layer.wq.data::class.simpleName}") + assertTrue(layer.ffnDown.data is NarrowFloatInputMajorTensorData) + assertContentEquals(values(dim * dim, Q_BASE), layer.wq.data.copyToFloatArray()) + assertTrue(layer.attnNorm.data is FloatArrayTensorData<*>, "F32 norms stay dense") + } + + private fun loader(policy: DTypePolicy) = DecoderSafeTensorsLoader( + ctx = DirectCpuExecutionContext(), + dtype = FP32::class, + metadata = metadata, + tiedEmbeddings = true, + dtypePolicy = policy, + ) + + private fun writeHfCheckpoint(): File { + val buffer = Buffer() + SafeTensorsWriter.write(buffer) { + tensorBF16("model.embed_tokens.weight", listOf(vocab.toLong(), dim.toLong()), values(vocab * dim, 1.0f)) + tensorF32("model.norm.weight", listOf(1L, dim.toLong()), values(dim, 2.0f)) + val l = "model.layers.0" + tensorF32("$l.input_layernorm.weight", listOf(dim.toLong()), values(dim, LN_BASE)) + tensorF32("$l.post_attention_layernorm.weight", listOf(dim.toLong()), values(dim, 2.5f)) + tensorBF16("$l.self_attn.q_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, Q_BASE)) + tensorBF16("$l.self_attn.k_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, 0.125f)) + tensorBF16("$l.self_attn.v_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, 0.25f)) + tensorBF16("$l.self_attn.o_proj.weight", listOf(dim.toLong(), dim.toLong()), values(dim * dim, 0.375f)) + tensorBF16("$l.mlp.gate_proj.weight", listOf(ffn.toLong(), dim.toLong()), values(ffn * dim, 0.5f)) + tensorBF16("$l.mlp.up_proj.weight", listOf(ffn.toLong(), dim.toLong()), values(ffn * dim, 0.625f)) + tensorBF16("$l.mlp.down_proj.weight", listOf(dim.toLong(), ffn.toLong()), values(dim * ffn, 0.75f)) + } + val file = Files.createTempFile("llama_st_fixture", ".safetensors").toFile().also { it.deleteOnExit() } + file.writeBytes(buffer.readByteArray()) + return file + } + + private fun values(n: Int, base: Float) = FloatArray(n) { base + (it % 16) * 0.125f } + + private companion object { + const val Q_BASE = -1.0f + const val LN_BASE = 0.5f + } +}