From e6280d1a09593a7a37b921c9f3a9639cf27519a4 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 15:08:18 +0200 Subject: [PATCH] =?UTF-8?q?feat(plan):=20footprint=20capability=20for=20sa?= =?UTF-8?q?fetensors=20and=20ONNX=20=E2=80=94=20header-only=20planInput,?= =?UTF-8?q?=20external=5Fdata=20fix,=20EDGE=20profile?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit There was no quick way to check whether a model fits on an embedded device with limited memory and compute before investing a day in converting it. This lands the library capability (#1169); the multi-format CLI over it incubates in SKaiNET-research until a stable core release carries these APIs. - StreamingSafeTensorsReader.planInput / StreamingOnnxReader.planInput: PlanInput with null geometry — weights-only plans through the same render/verdict/suggestion pipeline GGUF uses. Byte counts are authoritative: safetensors from data_offsets, ONNX from raw_data. - ONNX external_data (field 13) is parsed instead of skipped: >2 GB models keep their weights in a sibling file, and exactly those models previously reported ~0 bytes — a fit verdict that lied where it mattered most. Sizes are Long throughout (estimatedBytesLong; the Int view clamps instead of wrapping negative). - PlannerProfile.EDGE: an embedded device where the number the caller passes IS the usable RAM — reserve deliberately zero and documented to stay zero; weights mapped, KV auto-quantized past 80%. - SafeTensorsDataTypeMapper no longer printlns a WARNING into stdout mid-parse; UNKNOWN is the answer. - Tests pin the capability so later memory-layout refactors cannot silently lose it: per-format planInput sums, the 3 GiB external tensor (Long correctness), unknown-dtype pricing via offsets. Deferred (noted in #1169): HF config.json geometry, sharded index support, /proc/meminfo DeviceMemory provider, ONNX TensorId maps. Part of #1169. Co-Authored-By: Claude Fable 5 --- .../ROOT/pages/how-to/plan-model-memory.adoc | 28 ++++++ .../kotlin/sk/ainet/io/onnx/OnnxMemoryPlan.kt | 68 +++++++++++++ .../sk/ainet/io/onnx/StreamingOnnxReader.kt | 58 ++++++++--- .../sk/ainet/io/onnx/OnnxMemoryPlanTest.kt | 98 +++++++++++++++++++ .../safetensors/SafeTensorsDataTypeMapper.kt | 8 +- .../io/safetensors/SafeTensorsMemoryPlan.kt | 69 +++++++++++++ .../safetensors/SafeTensorsMemoryPlanTest.kt | 64 ++++++++++++ .../api/jvm/skainet-lang-core.api | 31 ++++++ .../ainet/lang/memory/plan/PlannerProfile.kt | 15 +++ 9 files changed, 423 insertions(+), 16 deletions(-) create mode 100644 skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxMemoryPlan.kt create mode 100644 skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxMemoryPlanTest.kt create mode 100644 skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlan.kt create mode 100644 skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlanTest.kt diff --git a/docs/modules/ROOT/pages/how-to/plan-model-memory.adoc b/docs/modules/ROOT/pages/how-to/plan-model-memory.adoc index 760f7c226..1caaa93f9 100644 --- a/docs/modules/ROOT/pages/how-to/plan-model-memory.adoc +++ b/docs/modules/ROOT/pages/how-to/plan-model-memory.adoc @@ -55,12 +55,40 @@ profiled.requireFits() // throws, naming the pool that ran out | `MOBILE_2GB` | 700 MB reserved, prefill chunked at 256, off-heap above 256 KB, KV auto-quantized to TurboQuant-4 once the plan passes 80 % of the budget, dispatcher dequantization warns above 5 % of bytes read, weights mapped | `DESKTOP` | The same reserve, no automatic KV quantization, heap staging | `NATIVE` | The smaller 300 MB Kotlin/Native reserve +| `EDGE` | An embedded device with limited memory and compute: **zero reserve** — the number you pass is the *usable* RAM, already net of what the OS holds — weights mapped, KV auto-quantized past 80 % | `PlannerProfile.forDevice(device)` | Picks `MOBILE_2GB` at or below 3 GB of RAM |=== A `ProfiledPlan` records what the profile decided — the KV switch appears as a note, not a silent rewrite — so a plan read months later says which rules produced it. +== Will it fit on an embedded device? Any format, in seconds + +The planner answers the pre-conversion question — *is it worth spending a day converting this +model for a device with ~2.1 GB of usable RAM?* — for **GGUF, safetensors and ONNX** files, from +the header/metadata only. A multi-gigabyte file is answered in seconds, and tensor payloads are +never read: + +[source,kotlin] +---- +val input = when (ModelFormat.fromFilePath(path)) { + ModelFormat.GGUF -> StreamingGGUFReader.open(src).planInput(ctx = 4096) + ModelFormat.SAFETENSORS -> StreamingSafeTensorsReader.open(src).planInput(modelName) + ModelFormat.ONNX -> StreamingOnnxReader.open(src).planInput(modelName) + null -> error("not a model file") +} +val verdict = PlannerProfile.EDGE.plan(input, availableBytes = parse("2.1G")) +println(verdict.render()) // ✔ fits / ✘ does not fit, with suggestions +---- + +`EDGE` treats the number you pass as the *usable* RAM — nothing is subtracted. ONNX weights kept +in a sibling `external_data` file are priced by their declared lengths, so multi-gigabyte models +report their real size. Two honest limits: safetensors and ONNX carry no architecture metadata, +so their plans are weights-only (`geometry == null` — KV cache and forward slab are modelled for +GGUF only); and file weights are a lower bound — leave headroom for the runtime's own buffers on +top of the verdict. A standalone multi-format CLI over this API is incubating in the +SKaiNET-research repository; the in-repo `skainet-plan` CLI covers GGUF. + == Check a real device, before allocating A phone has two memory pools, and one total cannot express the difference: the managed heap is diff --git a/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxMemoryPlan.kt b/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxMemoryPlan.kt new file mode 100644 index 000000000..03244444f --- /dev/null +++ b/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxMemoryPlan.kt @@ -0,0 +1,68 @@ +package sk.ainet.io.onnx + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.plan.PlanInput +import sk.ainet.lang.memory.plan.PlanTensor +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.FP64 +import sk.ainet.lang.types.Int16 +import sk.ainet.lang.types.Int32 +import sk.ainet.lang.types.Int64 +import sk.ainet.lang.types.Int8 + +/** + * Build a [PlanInput] from an ONNX model's **initializer table only** — the streaming reader + * records names, dims, dtypes and byte counts (including `external_data` lengths for >2 GB + * models) without materializing any tensor (#1169). + * + * ONNX carries no transformer-geometry metadata, so [PlanInput.geometry] is `null` and the plan + * is weights-only: KV cache and forward slab are not modelled, and the caller should say so. + */ +@ExperimentalMemoryApi +public fun StreamingOnnxReader.planInput( + modelName: String, + ctx: Int = 1, +): PlanInput { + val weights = tensors.map { t -> + PlanTensor( + name = t.name, + id = null, + format = onnxFormat(t.dataType, t.dataTypeName, t.estimatedBytesLong), + elementCount = t.nElements, + bytes = t.estimatedBytesLong, + ) + } + return PlanInput( + modelName = modelName, + architecture = "onnx", + weights = weights, + geometry = null, + ctx = ctx, + ) +} + +/** The [Format] an ONNX `TensorProto.DataType` describes; unknown types become an opaque encoding priced by [sizeInBytes]. */ +@ExperimentalMemoryApi +public fun onnxFormat(dataType: Int, dataTypeName: String, sizeInBytes: Long): Format { + val known: Pair? = when (dataType) { + 1 -> FP32 to 4 // FLOAT + 10 -> FP16 to 2 // FLOAT16 + 16 -> BF16 to 2 // BFLOAT16 + 11 -> FP64 to 8 // DOUBLE + 2, 3, 9 -> Int8 to 1 // UINT8, INT8, BOOL + 4, 5 -> Int16 to 2 // UINT16, INT16 + 6, 12 -> Int32 to 4 // INT32, UINT32 + 7, 13 -> Int64 to 8 // INT64, UINT64 + else -> null + } + return if (known != null) { + Format(known.first, TensorEncoding.Dense(known.second)) + } else { + Format(FP32, TensorEncoding.Opaque(dataTypeName, sizeInBytes)) + } +} diff --git a/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/StreamingOnnxReader.kt b/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/StreamingOnnxReader.kt index 8c95e0b68..bcd5a9c62 100644 --- a/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/StreamingOnnxReader.kt +++ b/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/StreamingOnnxReader.kt @@ -210,6 +210,9 @@ public class StreamingOnnxReader private constructor( var rawDataOffset = -1L var rawDataLength = 0 var hasTypedData = false + var externalLocation: String? = null + var externalOffset = -1L + var externalLength = -1L while (reader.hasRemaining(endPos)) { val tag = reader.readVarint() @@ -259,8 +262,27 @@ public class StreamingOnnxReader private constructor( reader.skipField(wireType) } 13 -> { - // external_data (repeated) - indicates external storage - reader.skipField(wireType) + // external_data (repeated StringStringEntryProto): weights of >2 GB models + // live in a sibling file, described by key/value pairs. Without parsing + // these, exactly the models that do not fit reported ~0 bytes (#1169). + val entryLength = reader.readVarint().toInt() + val entryEnd = reader.position + entryLength + var key = "" + var value = "" + while (reader.hasRemaining(entryEnd)) { + val entryTag = reader.readVarint() + when (ProtobufWireReader.fieldNumber(entryTag)) { + 1 -> key = reader.readString() + 2 -> value = reader.readString() + else -> reader.skipField(ProtobufWireReader.wireType(entryTag)) + } + } + when (key) { + "location" -> externalLocation = value + "offset" -> externalOffset = value.toLongOrNull() ?: -1L + "length" -> externalLength = value.toLongOrNull() ?: -1L + } + reader.seek(entryEnd) } else -> reader.skipField(wireType) } @@ -269,12 +291,12 @@ public class StreamingOnnxReader private constructor( if (name.isNotEmpty()) { val nElements = if (dims.isEmpty()) 0L else dims.fold(1L) { acc, d -> acc * d } val typeSize = getDataTypeSize(dataType) - val estimatedBytes = if (rawDataLength > 0) { - rawDataLength - } else if (hasTypedData && nElements > 0 && typeSize > 0) { - (nElements * typeSize).toInt() - } else { - 0 + // Long throughout: a single >2 GB initializer must not wrap (#1169). + val estimatedBytesLong: Long = when { + rawDataLength > 0 -> rawDataLength.toLong() + externalLength > 0 -> externalLength + nElements > 0 && typeSize > 0 -> nElements * typeSize + else -> 0L } _tensors.add( @@ -286,8 +308,12 @@ public class StreamingOnnxReader private constructor( nElements = nElements, rawDataOffset = rawDataOffset, rawDataLength = rawDataLength, - estimatedBytes = estimatedBytes, - hasTypedArrayData = hasTypedData && rawDataLength <= 0 + estimatedBytes = estimatedBytesLong.coerceAtMost(Int.MAX_VALUE.toLong()).toInt(), + hasTypedArrayData = hasTypedData && rawDataLength <= 0, + externalLocation = externalLocation, + externalOffset = externalOffset, + externalLength = externalLength, + estimatedBytesLong = estimatedBytesLong, ) ) } @@ -359,8 +385,16 @@ public data class StreamingOnnxTensorInfo( val rawDataOffset: Long, /** Length of raw_data in bytes (0 if not available) */ val rawDataLength: Int, - /** Estimated size in bytes (from raw_data or calculated) */ + /** Estimated size in bytes (from raw_data or calculated), clamped to Int.MAX_VALUE — prefer [estimatedBytesLong] */ val estimatedBytes: Int, /** True if tensor data is in typed arrays (requires full parsing) */ - val hasTypedArrayData: Boolean + val hasTypedArrayData: Boolean, + /** external_data `location` (a sibling file path), or null when the data is in this file */ + val externalLocation: String? = null, + /** external_data `offset` within [externalLocation], -1 when absent */ + val externalOffset: Long = -1L, + /** external_data `length` in bytes, -1 when absent */ + val externalLength: Long = -1L, + /** Estimated size in bytes without the Int clamp — raw_data, external length, or elements × type size */ + val estimatedBytesLong: Long = estimatedBytes.toLong() ) diff --git a/skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxMemoryPlanTest.kt b/skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxMemoryPlanTest.kt new file mode 100644 index 000000000..dcc144d60 --- /dev/null +++ b/skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxMemoryPlanTest.kt @@ -0,0 +1,98 @@ +package sk.ainet.io.onnx + +import onnx.GraphProto +import onnx.ModelProto +import onnx.StringStringEntryProto +import onnx.TensorProto +import pbandk.ByteArr +import pbandk.encodeToByteArray +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import java.nio.file.Files +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** + * #1169: footprint planning from an ONNX initializer table — including `external_data`, whose + * lengths previously reported ~0 bytes for exactly the >2 GB models that do not fit. + */ +@OptIn(ExperimentalMemoryApi::class) +class OnnxMemoryPlanTest { + + private fun model(vararg tensors: TensorProto): ByteArray = + ModelProto(irVersion = 8, producerName = "OnnxMemoryPlanTest", graph = GraphProto(name = "g", initializer = tensors.toList())) + .encodeToByteArray() + + private fun withReader(bytes: ByteArray, block: (StreamingOnnxReader) -> R): R { + val f = Files.createTempFile("plan_model", ".onnx").toFile() + f.deleteOnExit() + f.writeBytes(bytes) + return StreamingOnnxReader.open(JvmRandomAccessSource.open(f)).use(block) + } + + @Test + fun planInputPricesRawDataInitializers() { + val bytes = model( + TensorProto(name = "w", dims = listOf(4L, 8L), dataType = TensorProto.DataType.FLOAT.value, rawData = ByteArr(ByteArray(128))), + TensorProto(name = "b", dims = listOf(8L), dataType = TensorProto.DataType.FLOAT.value, rawData = ByteArr(ByteArray(32))), + ) + withReader(bytes) { reader -> + val input = reader.planInput(modelName = "test.onnx") + assertEquals(2, input.weights.size) + assertNull(input.geometry, "ONNX carries no architecture metadata") + assertEquals("onnx", input.architecture) + val w = input.weights.first { it.name == "w" } + assertEquals(32L, w.elementCount) + assertEquals(128L, w.bytes) + assertEquals(160L, input.weights.sumOf { it.bytes }) + } + } + + @Test + fun externalDataLengthIsPricedNotZero() { + // A weight stored in a sibling file: no raw_data, external_data carries location/offset/length. + val threeGiB = 3L * 1024 * 1024 * 1024 + val bytes = model( + TensorProto( + name = "big", + dims = listOf(threeGiB / 4), + dataType = TensorProto.DataType.FLOAT.value, + externalData = listOf( + StringStringEntryProto(key = "location", value = "model.onnx_data"), + StringStringEntryProto(key = "offset", value = "0"), + StringStringEntryProto(key = "length", value = threeGiB.toString()), + ), + dataLocation = TensorProto.DataLocation.EXTERNAL, + ), + ) + withReader(bytes) { reader -> + val t = reader.tensors.single() + assertEquals("model.onnx_data", t.externalLocation) + assertEquals(threeGiB, t.externalLength) + assertEquals(threeGiB, t.estimatedBytesLong, "external length must be priced, in Long") + assertEquals(Int.MAX_VALUE, t.estimatedBytes, "Int view clamps instead of wrapping negative") + val input = reader.planInput(modelName = "big.onnx") + assertEquals(threeGiB, input.weights.single().bytes) + } + } + + @Test + fun externalDataWithoutLengthFallsBackToElementsTimesTypeSize() { + val bytes = model( + TensorProto( + name = "ext", + dims = listOf(1024L), + dataType = TensorProto.DataType.FLOAT.value, + externalData = listOf(StringStringEntryProto(key = "location", value = "model.onnx_data")), + dataLocation = TensorProto.DataLocation.EXTERNAL, + ), + ) + withReader(bytes) { reader -> + val t = reader.tensors.single() + assertEquals(4096L, t.estimatedBytesLong, "no length entry: elements × type size") + assertTrue(t.externalLocation != null) + } + } +} diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsDataTypeMapper.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsDataTypeMapper.kt index 8b1865a65..1d693e2b5 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsDataTypeMapper.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsDataTypeMapper.kt @@ -32,10 +32,10 @@ object SafeTensorsDataTypeMapper { "F64" -> DataType.FLOAT64 "Q4" -> DataType.QUANT4 "Q8" -> DataType.QUANT8 - else -> { - println("WARNING: Unknown SafeTensors dtype: $safeTensorsType") - DataType.UNKNOWN - } + // No println here: this runs inside header parsing, and a diagnostic on stdout would + // corrupt machine-readable output (e.g. the skainet-plan table, #1169). UNKNOWN is + // the answer; callers who care can check for it. + else -> DataType.UNKNOWN } } diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlan.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlan.kt new file mode 100644 index 000000000..937c86eb1 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlan.kt @@ -0,0 +1,69 @@ +package sk.ainet.io.safetensors + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.plan.PlanInput +import sk.ainet.lang.memory.plan.PlanTensor +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.FP64 +import sk.ainet.lang.types.Int16 +import sk.ainet.lang.types.Int32 +import sk.ainet.lang.types.Int64 +import sk.ainet.lang.types.Int8 + +/** + * Build a [PlanInput] from a safetensors **header only** — the JSON tensor table at the start of + * the file; no tensor payload is read (#1169). + * + * safetensors carries no architecture metadata, so [PlanInput.geometry] is `null` and the plan is + * weights-only: KV cache and forward slab are not modelled, and the caller should say so. + * Per-tensor byte counts come from the header's `data_offsets` — authoritative even for dtypes + * with no fixed per-element width. + */ +@ExperimentalMemoryApi +public fun StreamingSafeTensorsReader.planInput( + modelName: String, + ctx: Int = 1, +): PlanInput { + val weights = tensors.map { t -> + PlanTensor( + name = t.name, + id = null, + format = safeTensorsFormat(t.dtype, t.sizeInBytes), + elementCount = t.elementCount, + bytes = t.sizeInBytes, + ) + } + return PlanInput( + modelName = modelName, + architecture = "safetensors", + weights = weights, + geometry = null, + ctx = ctx, + ) +} + +/** The [Format] a safetensors dtype string describes; unknown dtypes become an opaque encoding priced by [sizeInBytes]. */ +@ExperimentalMemoryApi +public fun safeTensorsFormat(dtype: String, sizeInBytes: Long): Format { + val known: Pair? = when (dtype) { + "F32" -> FP32 to 4 + "F16" -> FP16 to 2 + "BF16" -> BF16 to 2 + "F64" -> FP64 to 8 + "I8", "U8", "BOOL" -> Int8 to 1 + "I16", "U16" -> Int16 to 2 + "I32", "U32" -> Int32 to 4 + "I64", "U64" -> Int64 to 8 + else -> null + } + return if (known != null) { + Format(known.first, TensorEncoding.Dense(known.second)) + } else { + Format(FP32, TensorEncoding.Opaque(dtype, sizeInBytes)) + } +} diff --git a/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlanTest.kt b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlanTest.kt new file mode 100644 index 000000000..f831ed29b --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsMemoryPlanTest.kt @@ -0,0 +1,64 @@ +package sk.ainet.io.safetensors + +import sk.ainet.io.RandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull + +/** + * #1169: footprint planning from a safetensors header — sizes come from `data_offsets`, so they + * are authoritative even for dtypes with no fixed per-element width. + */ +@OptIn(ExperimentalMemoryApi::class) +class SafeTensorsMemoryPlanTest { + + private class ByteSource(private val data: ByteArray) : RandomAccessSource { + override val size: Long = data.size.toLong() + override fun readAt(position: Long, length: Int): ByteArray = + data.copyOfRange(position.toInt(), (position + length).toInt()) + override fun readAt(position: Long, buffer: ByteArray, offset: Int, length: Int): Int { + val n = minOf(length, (size - position).toInt()) + data.copyInto(buffer, offset, position.toInt(), position.toInt() + n) + return n + } + override fun close() {} + } + + /** 8-byte LE header length + JSON header + zero payload — the whole format. */ + private fun file(headerJson: String, payloadBytes: Int): ByteArray { + val header = headerJson.encodeToByteArray() + val out = ByteArray(8 + header.size + payloadBytes) + var len = header.size.toLong() + for (i in 0 until 8) { out[i] = (len and 0xFF).toByte(); len = len shr 8 } + header.copyInto(out, 8) + return out + } + + @Test + fun planInputPricesFromDataOffsets() { + val json = """{"w":{"dtype":"F32","shape":[4,8],"data_offsets":[0,128]},""" + + """"h":{"dtype":"BF16","shape":[8],"data_offsets":[128,144]}}""" + val reader = StreamingSafeTensorsReader.open(ByteSource(file(json, 144))) + val input = reader.planInput(modelName = "test.safetensors") + assertEquals(2, input.weights.size) + assertNull(input.geometry, "safetensors carries no architecture metadata") + assertEquals("safetensors", input.architecture) + val w = input.weights.first { it.name == "w" } + assertEquals(32L, w.elementCount) + assertEquals(128L, w.bytes) + assertEquals(w.bytes, w.residentBytes, "no form resolved: resident = as stored") + assertEquals(144L, input.weights.sumOf { it.bytes }) + } + + @Test + fun unknownDtypeIsPricedByItsOffsets() { + // A quantized/unknown dtype has no per-element width — data_offsets still price it. + val json = """{"q":{"dtype":"Q4_K_M","shape":[256],"data_offsets":[0,144]}}""" + val reader = StreamingSafeTensorsReader.open(ByteSource(file(json, 144))) + val input = reader.planInput(modelName = "q.safetensors") + val q = input.weights.single() + assertEquals(144L, q.bytes) + assertEquals("Q4_K_M", q.format.encoding.name) + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 11566d792..481b89b8a 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -1807,6 +1807,7 @@ public final class sk/ainet/lang/memory/plan/PlannerProfile { public final class sk/ainet/lang/memory/plan/PlannerProfile$Companion { public final fun forDevice (Lsk/ainet/lang/memory/plan/DeviceMemory;)Lsk/ainet/lang/memory/plan/PlannerProfile; public final fun getDESKTOP ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun getEDGE ()Lsk/ainet/lang/memory/plan/PlannerProfile; public final fun getMOBILE_2GB ()Lsk/ainet/lang/memory/plan/PlannerProfile; public final fun getNATIVE ()Lsk/ainet/lang/memory/plan/PlannerProfile; } @@ -4659,6 +4660,36 @@ public final class sk/ainet/lang/tensor/data/Bf16TensorDataKt { public static final fun toFloatArray (Lsk/ainet/lang/tensor/data/Bf16TensorData;)[F } +public final class sk/ainet/lang/tensor/data/BitNetB158TensorData : sk/ainet/lang/tensor/data/TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { + public static final field Companion Lsk/ainet/lang/tensor/data/BitNetB158TensorData$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun copyToFloatArray ()[F + public fun dequantizeBlock (I[FI)V + public fun get ([I)Ljava/lang/Byte; + public synthetic fun get ([I)Ljava/lang/Object; + public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; + public fun getBlockSize ()I + public fun getElementCount ()J + public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public fun getPackedData ()[B + public fun getPackedView ()Lsk/ainet/lang/memory/TensorView; + public fun getPhysicalBytes ()J + public final fun getScale ()F + public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; + public fun set ([IB)V + public synthetic fun set ([ILjava/lang/Object;)V + public fun toFloatArray ()[F + public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; +} + +public final class sk/ainet/lang/tensor/data/BitNetB158TensorData$Companion { + public final fun fromFloats (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/data/BitNetB158TensorData; + public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/BitNetB158TensorData; +} + public final class sk/ainet/lang/tensor/data/DenseFloatArrayTensorData : sk/ainet/lang/tensor/data/FloatArrayTensorData { public fun (Lsk/ainet/lang/tensor/Shape;[F)V public fun copyToFloatArray ()[F diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt index f62a4c001..bf75d8e4b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt @@ -125,6 +125,21 @@ public data class PlannerProfile( strict = true, ) + /** + * An embedded device with limited memory and compute (#1169): the number the caller passes + * as available bytes IS the usable RAM — already net of what the OS holds — so + * [reserveBytes] is deliberately **zero and must stay zero**; adding a reserve here would + * double-count what the caller already subtracted. Weights are counted mapped, the KV + * cache auto-quantizes past 80 % of the budget, and the profile is not strict: this + * profile's job is a pre-flight verdict, not a refusal. + */ + public val EDGE: PlannerProfile = PlannerProfile( + name = "edge", + reserveBytes = 0, + kvAutoQuantizeAbove = 0.80, + weightsMapped = true, + ) + /** A desktop or server JVM: the same reserve, no automatic KV quantization, heap staging. */ public val DESKTOP: PlannerProfile = PlannerProfile( name = "desktop",