From f0e2e08984904e450142b4c23fe88d057b84f100 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sat, 22 Aug 2026 22:37:27 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20MemoryPlan=20from=20the=20GGUF?= =?UTF-8?q?=20header=20=E2=80=94=20weights,=20KV=20(bf16/TurboQuant),=20fo?= =?UTF-8?q?rward=20slab,=20headroom,=20budget=20fit,=20suggestions=20(SKEE?= =?UTF-8?q?P-003=20P1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M0 (#1001), PRD M0-F1..F3 / M0-A4: "will this model fit on this device at this context length?" answered from shapes and encodings only — no tensor bytes are read. - sk.ainet.lang.memory.plan: PlanTensor (name, TensorId?, Format, count, bytes; allocation = AllocationSpec mapped/MODEL/read-only), ModelGeometry, KvCacheMode (BF16, TURBOQUANT_4 via TensorEncoding.TurboQuantPolar), PlanInput (model, weights, geometry, ctx, prefill chunk 256, kv mode; unmappedWeights listed, never dropped), Budget (explicit or available − reserve: 700 MB Android/JVM, 300 MB native — decision #11), MemoryPlan (weights resident · kv @ ctx with the alternate mode · forward slab · heap headroom · total · fits · ≥ 2 suggestions with savings: --kv turboquant, --ctx N/2, smaller model; render() = the PRD §4.3 table), MemoryPlans.plan(input, budget) with the documented estimates. - io-gguf: StreamingGGUFReader.planInput(ctx, prefillChunk, kvMode, nameMap) — header only (tensor table + .block_count / embedding_length / attention.head_count(_kv) / key_length / value_length / feed_forward_length / vocab_size / context_length); ggufGeometry(); ggufFormat(type, nBytes). - Tests: MemoryPlanTest (Llama-3.2-1B-like geometry: KV bf16 @2048 = 64 MiB, TurboQuant ≈ ¼, forward slab scaling, totals/residency, does-not-fit suggestions, TurboQuant mode + available-memory budget, no geometry, unmapped, byte formatting); GgufMemoryPlanTest (synthetic header-only plan, format mapping, fixture-gated Qwen2.5-0.5B real header). - BCV: lang-core jvm dump regenerated (additions only). Includes the AllocationSpec commit of #1053 (cherry-picked: the plan's line items are AllocationSpecs); that commit drops out on rebase once Closes #1012 Co-Authored-By: Claude Fable 5 --- .../kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt | 85 +++++++ .../sk/ainet/io/gguf/GgufMemoryPlanTest.kt | 72 ++++++ .../api/jvm/skainet-lang-core.api | 179 +++++++++++++ .../sk/ainet/lang/memory/plan/MemoryPlan.kt | 237 ++++++++++++++++++ .../ainet/lang/memory/plan/MemoryPlanTest.kt | 127 ++++++++++ 5 files changed, 700 insertions(+) create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufMemoryPlanTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt new file mode 100644 index 000000000..6c7f9ee9a --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt @@ -0,0 +1,85 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.weights.NameMap +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.plan.KvCacheMode +import sk.ainet.lang.memory.plan.ModelGeometry +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 GGUF **header only** — tensor names, shapes and types plus the + * architecture metadata keys; no tensor bytes are read (PRD M0-F1). [ctx] defaults to the trained + * context length of the model when the header has one. + */ +@ExperimentalMemoryApi +public fun StreamingGGUFReader.planInput( + ctx: Int? = null, + prefillChunk: Int = PlanInput.DEFAULT_PREFILL_CHUNK, + kvMode: KvCacheMode = KvCacheMode.BF16, + nameMap: NameMap? = nameMap(), +): PlanInput { + val arch = fields["general.architecture"] as? String ?: "unknown" + val name = fields["general.name"] as? String ?: arch + val geometry = ggufGeometry(arch) + val weights = tensors.map { t -> + val format = ggufFormat(t.tensorType, t.nBytes) + PlanTensor( + name = t.name, + id = nameMap?.toTensorId(t.name), + format = format, + elementCount = t.nElements, + bytes = format.physicalBytes(t.nElements) ?: t.nBytes, + ) + } + val ctxUsed = ctx ?: geometry?.trainedContextLength ?: 2048 + return PlanInput(name, arch, weights, geometry, ctxUsed, prefillChunk, kvMode) +} + +/** Architecture metadata (`.block_count`, `.attention.head_count`, …) as a [ModelGeometry], or `null` if the header lacks it. */ +@ExperimentalMemoryApi +public fun StreamingGGUFReader.ggufGeometry(architecture: String? = fields["general.architecture"] as? String): ModelGeometry? { + val arch = architecture ?: return null + fun int(key: String): Int? = (fields["$arch.$key"] as? Number)?.toInt() ?: (fields["$arch.$key"] as? UInt)?.toInt() + val layers = int("block_count") ?: return null + val emb = int("embedding_length") ?: return null + val heads = int("attention.head_count") ?: return null + val kvHeads = int("attention.head_count_kv") ?: heads + val headDim = int("attention.key_length") ?: (emb / heads) + val valueDim = int("attention.value_length") ?: headDim + val ffn = int("feed_forward_length") ?: (4 * emb) + val vocab = int("vocab_size") ?: (fields["tokenizer.ggml.tokens"] as? List<*>)?.size + ?: tensors.firstOrNull { it.name == "token_embd.weight" }?.shape?.lastOrNull()?.toInt() ?: 0 + return ModelGeometry(layers, heads, kvHeads, headDim, valueDim, emb, ffn, vocab, int("context_length")) +} + +/** `Format` of a GGUF tensor type: quantized types are logically FP32 with their block encoding. */ +@ExperimentalMemoryApi +public fun ggufFormat(type: GGMLQuantizationType, nBytes: Long): Format { + val dtype: DType = when (type) { + GGMLQuantizationType.F32 -> FP32; GGMLQuantizationType.F16 -> FP16; GGMLQuantizationType.BF16 -> BF16; GGMLQuantizationType.F64 -> FP64 + GGMLQuantizationType.I8 -> Int8; GGMLQuantizationType.I16 -> Int16; GGMLQuantizationType.I32 -> Int32; GGMLQuantizationType.I64 -> Int64 + else -> FP32 + } + val encoding: TensorEncoding = when (type) { + GGMLQuantizationType.F32 -> TensorEncoding.Dense(4); GGMLQuantizationType.F16, GGMLQuantizationType.BF16 -> TensorEncoding.Dense(2) + GGMLQuantizationType.F64 -> TensorEncoding.Dense(8); GGMLQuantizationType.I8 -> TensorEncoding.Dense(1) + GGMLQuantizationType.I16 -> TensorEncoding.Dense(2); GGMLQuantizationType.I32 -> TensorEncoding.Dense(4); GGMLQuantizationType.I64 -> TensorEncoding.Dense(8) + GGMLQuantizationType.Q4_0 -> TensorEncoding.Q4_0; GGMLQuantizationType.Q5_0 -> TensorEncoding.Q5_0; GGMLQuantizationType.Q5_1 -> TensorEncoding.Q5_1 + GGMLQuantizationType.Q8_0 -> TensorEncoding.Q8_0; GGMLQuantizationType.Q4_K -> TensorEncoding.Q4_K; GGMLQuantizationType.Q5_K -> TensorEncoding.Q5_K + GGMLQuantizationType.Q6_K -> TensorEncoding.Q6_K + else -> TensorEncoding.Opaque(type.name, nBytes) + } + return Format(dtype, encoding) +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufMemoryPlanTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufMemoryPlanTest.kt new file mode 100644 index 000000000..6fc29ad11 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufMemoryPlanTest.kt @@ -0,0 +1,72 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.Budget +import sk.ainet.lang.memory.plan.MemoryPlans +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** SKEEP-003 M0-F1: the plan comes from the GGUF header alone — tensor table + metadata, no tensor bytes. */ +@OptIn(ExperimentalMemoryApi::class) +class GgufMemoryPlanTest { + + @Test + fun planFromSyntheticHeaderOnly() { + val q4 = SyntheticGguf.tensor("w.q4k", GGMLQuantizationType.Q4_K, 512) + val q8 = SyntheticGguf.tensor("w.q8", GGMLQuantizationType.Q8_0, 64) + val f32 = SyntheticGguf.tensor("w.f32", GGMLQuantizationType.F32, 10) + val file = SyntheticGguf.write(q4, q8, f32) + JvmRandomAccessSource.open(file).use { src -> + val reader = StreamingGGUFReader.open(src) + val input = reader.planInput(ctx = 256) + assertEquals("test", input.architecture) + assertNull(input.geometry) // synthetic header carries no architecture geometry + assertEquals(3, input.weights.size) + val byName = input.weights.associateBy { it.name } + assertEquals(FP32, byName["w.q4k"]!!.format.dtype); assertEquals(TensorEncoding.Q4_K, byName["w.q4k"]!!.format.encoding) + assertEquals(144L * 2, byName["w.q4k"]!!.bytes) + assertEquals(34L * 2, byName["w.q8"]!!.bytes) + assertEquals(TensorEncoding.Dense(4), byName["w.f32"]!!.format.encoding); assertEquals(40L, byName["w.f32"]!!.bytes) + // unknown architecture: no name map, names kept as unmapped + assertEquals(listOf("w.q4k", "w.q8", "w.f32"), input.unmappedWeights) + val plan = MemoryPlans.plan(input, Budget.of(1L shl 30)) + assertEquals(288L + 68 + 40, plan.weightsBytes) + assertEquals(true, plan.fits) + } + } + + @Test + fun ggufFormatMapping() { + assertEquals(TensorEncoding.Q6_K, ggufFormat(GGMLQuantizationType.Q6_K, 0).encoding) + assertEquals(FP32, ggufFormat(GGMLQuantizationType.Q6_K, 0).dtype) + assertEquals(TensorEncoding.Dense(2), ggufFormat(GGMLQuantizationType.BF16, 0).encoding) + assertTrue(ggufFormat(GGMLQuantizationType.TQ1_0, 123).encoding is TensorEncoding.Opaque) + } + + /** Real file, fixture-gated (see GgufNameMapFixtureTest): the plan must be consistent with the header. */ + @Test + fun planFromQwenFixture() { + val dir = File(System.getProperty("skainet.test.fixturesDir") ?: "../skainet-io-core/build/test-fixtures") + val f = File(dir, "Qwen2.5-0.5B-Instruct-Q8_0.gguf") + if (!f.isFile) { println("[skip] ${f.name} not present"); return } + JvmRandomAccessSource.open(f).use { src -> + val reader = StreamingGGUFReader.open(src) + val input = reader.planInput(ctx = 2048) + val g = assertNotNull(input.geometry) + assertEquals(24, g.layers); assertEquals(896, g.embeddingLength); assertEquals(14, g.heads); assertEquals(2, g.kvHeads) + assertTrue(input.unmappedWeights.isEmpty(), "unmapped: ${input.unmappedWeights}") + val plan = MemoryPlans.plan(input, Budget.of(1300L shl 20)) + // Q8_0 0.5B: weights ≈ file size (header excluded), within 2 % + val packedBytes = reader.tensors.sumOf { it.nBytes } + assertTrue(kotlin.math.abs(plan.weightsBytes - packedBytes) <= packedBytes / 50, "weights ${plan.weightsBytes} vs packed $packedBytes") + println(plan.render()) + } + } +} 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 80757e161..5d22ae6a7 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 @@ -709,6 +709,185 @@ public final class sk/ainet/lang/memory/ScopeKind : java/lang/Enum { public static fun values ()[Lsk/ainet/lang/memory/ScopeKind; } +public final class sk/ainet/lang/memory/plan/Budget { + public static final field Companion Lsk/ainet/lang/memory/plan/Budget$Companion; + public static final field RESERVE_ANDROID_JVM J + public static final field RESERVE_NATIVE J + public fun (JLjava/lang/String;)V + public final fun component1 ()J + public final fun component2 ()Ljava/lang/String; + public final fun copy (JLjava/lang/String;)Lsk/ainet/lang/memory/plan/Budget; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/Budget;JLjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/Budget; + public fun equals (Ljava/lang/Object;)Z + public final fun getBytes ()J + public final fun getDescription ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/Budget$Companion { + public final fun available (JJ)Lsk/ainet/lang/memory/plan/Budget; + public static synthetic fun available$default (Lsk/ainet/lang/memory/plan/Budget$Companion;JJILjava/lang/Object;)Lsk/ainet/lang/memory/plan/Budget; + public final fun of (J)Lsk/ainet/lang/memory/plan/Budget; +} + +public final class sk/ainet/lang/memory/plan/KvCacheMode : java/lang/Enum { + public static final field BF16 Lsk/ainet/lang/memory/plan/KvCacheMode; + public static final field TURBOQUANT_4 Lsk/ainet/lang/memory/plan/KvCacheMode; + public final fun bytes (J)J + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public final fun getLabel ()Ljava/lang/String; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/memory/plan/KvCacheMode; + public static fun values ()[Lsk/ainet/lang/memory/plan/KvCacheMode; +} + +public final class sk/ainet/lang/memory/plan/MemoryPlan { + public fun (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;)V + public final fun component1 ()Lsk/ainet/lang/memory/plan/PlanInput; + public final fun component2 ()J + public final fun component3 ()J + public final fun component4 ()J + public final fun component5 ()J + public final fun component6 ()J + public final fun component7 ()Lsk/ainet/lang/memory/plan/Budget; + public final fun copy (Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;)Lsk/ainet/lang/memory/plan/MemoryPlan; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/PlanInput;JJJJJLsk/ainet/lang/memory/plan/Budget;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/MemoryPlan; + public fun equals (Ljava/lang/Object;)Z + public final fun getBudget ()Lsk/ainet/lang/memory/plan/Budget; + public final fun getFits ()Ljava/lang/Boolean; + public final fun getForwardBytes ()J + public final fun getHeadroomBytes ()J + public final fun getInput ()Lsk/ainet/lang/memory/plan/PlanInput; + public final fun getKvBytes ()J + public final fun getKvBytesAlternate ()J + public final fun getLines ()Ljava/util/List; + public final fun getResidentBytes ()J + public final fun getTotalBytes ()J + public final fun getWeightsBytes ()J + public fun hashCode ()I + public final fun render ()Ljava/lang/String; + public final fun suggestions ()Ljava/util/List; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/MemoryPlans { + public static final field HEAP_HEADROOM_BYTES J + public static final field INSTANCE Lsk/ainet/lang/memory/plan/MemoryPlans; + public final fun formatBytes (J)Ljava/lang/String; + public final fun forwardBytes (Lsk/ainet/lang/memory/plan/ModelGeometry;II)J + public final fun kvElements (Lsk/ainet/lang/memory/plan/ModelGeometry;I)J + public final fun plan (Lsk/ainet/lang/memory/plan/PlanInput;Lsk/ainet/lang/memory/plan/Budget;)Lsk/ainet/lang/memory/plan/MemoryPlan; + public static synthetic fun plan$default (Lsk/ainet/lang/memory/plan/MemoryPlans;Lsk/ainet/lang/memory/plan/PlanInput;Lsk/ainet/lang/memory/plan/Budget;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/MemoryPlan; +} + +public final class sk/ainet/lang/memory/plan/ModelGeometry { + public fun (IIIIIIIILjava/lang/Integer;)V + public synthetic fun (IIIIIIIILjava/lang/Integer;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()I + public final fun component2 ()I + public final fun component3 ()I + public final fun component4 ()I + public final fun component5 ()I + public final fun component6 ()I + public final fun component7 ()I + public final fun component8 ()I + public final fun component9 ()Ljava/lang/Integer; + public final fun copy (IIIIIIIILjava/lang/Integer;)Lsk/ainet/lang/memory/plan/ModelGeometry; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/ModelGeometry;IIIIIIIILjava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/ModelGeometry; + public fun equals (Ljava/lang/Object;)Z + public final fun getEmbeddingLength ()I + public final fun getFeedForwardLength ()I + public final fun getHeadDim ()I + public final fun getHeads ()I + public final fun getKvHeads ()I + public final fun getLayers ()I + public final fun getTrainedContextLength ()Ljava/lang/Integer; + public final fun getValueDim ()I + public final fun getVocabSize ()I + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/PlanInput { + public static final field Companion Lsk/ainet/lang/memory/plan/PlanInput$Companion; + public static final field DEFAULT_PREFILL_CHUNK I + public fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/memory/plan/ModelGeometry;IILsk/ainet/lang/memory/plan/KvCacheMode;)V + public synthetic fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/memory/plan/ModelGeometry;IILsk/ainet/lang/memory/plan/KvCacheMode;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()Ljava/util/List; + public final fun component4 ()Lsk/ainet/lang/memory/plan/ModelGeometry; + public final fun component5 ()I + public final fun component6 ()I + public final fun component7 ()Lsk/ainet/lang/memory/plan/KvCacheMode; + public final fun copy (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/memory/plan/ModelGeometry;IILsk/ainet/lang/memory/plan/KvCacheMode;)Lsk/ainet/lang/memory/plan/PlanInput; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanInput;Ljava/lang/String;Ljava/lang/String;Ljava/util/List;Lsk/ainet/lang/memory/plan/ModelGeometry;IILsk/ainet/lang/memory/plan/KvCacheMode;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanInput; + public fun equals (Ljava/lang/Object;)Z + public final fun getArchitecture ()Ljava/lang/String; + public final fun getCtx ()I + public final fun getGeometry ()Lsk/ainet/lang/memory/plan/ModelGeometry; + public final fun getKvMode ()Lsk/ainet/lang/memory/plan/KvCacheMode; + public final fun getModelName ()Ljava/lang/String; + public final fun getPrefillChunk ()I + public final fun getUnmappedWeights ()Ljava/util/List; + public final fun getWeights ()Ljava/util/List; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/PlanInput$Companion { +} + +public final class sk/ainet/lang/memory/plan/PlanLine { + public fun (Ljava/lang/String;Ljava/lang/String;JZ)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()J + public final fun component4 ()Z + public final fun copy (Ljava/lang/String;Ljava/lang/String;JZ)Lsk/ainet/lang/memory/plan/PlanLine; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanLine;Ljava/lang/String;Ljava/lang/String;JZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanLine; + public fun equals (Ljava/lang/Object;)Z + public final fun getBytes ()J + public final fun getDetail ()Ljava/lang/String; + public final fun getResident ()Z + public final fun getSection ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/PlanTensor { + public fun (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJ)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Lsk/ainet/lang/tensor/TensorId; + public final fun component3 ()Lsk/ainet/lang/memory/Format; + public final fun component4 ()J + public final fun component5 ()J + public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJ)Lsk/ainet/lang/memory/plan/PlanTensor; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanTensor;Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanTensor; + public fun equals (Ljava/lang/Object;)Z + public final fun getAllocation ()Lsk/ainet/lang/memory/AllocationSpec; + public final fun getBytes ()J + public final fun getElementCount ()J + public final fun getFormat ()Lsk/ainet/lang/memory/Format; + public final fun getId ()Lsk/ainet/lang/tensor/TensorId; + public final fun getName ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/Suggestion { + public fun (Ljava/lang/String;J)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()J + public final fun copy (Ljava/lang/String;J)Lsk/ainet/lang/memory/plan/Suggestion; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/Suggestion;Ljava/lang/String;JILjava/lang/Object;)Lsk/ainet/lang/memory/plan/Suggestion; + public fun equals (Ljava/lang/Object;)Z + public final fun getSavesBytes ()J + public final fun getText ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/nn/AvgPool2d : sk/ainet/lang/nn/Module { public fun (Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;)V public synthetic fun (Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt new file mode 100644 index 000000000..7084cb487 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt @@ -0,0 +1,237 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.AllocationSpec +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ScopeKind +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.MemoryDomain +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * One weight tensor as the planner sees it — shape and format only, never bytes (SKEEP-003 §8 item 1, + * PRD M0-F1). [id] is the [TensorId] the checkpoint's `NameMap` produced, `null` when unmapped. + */ +@ExperimentalMemoryApi +public data class PlanTensor( + val name: String, + val id: TensorId?, + val format: Format, + val elementCount: Long, + /** Physical bytes; falls back to the checkpoint's own byte count when the encoding cannot compute it. */ + val bytes: Long, +) { + /** The allocation this weight needs: mapped, model-lifetime, read-only. */ + val allocation: AllocationSpec + get() = AllocationSpec(format, elementCount, MemoryDomain.MMAP_FILE, ScopeKind.MODEL, mutable = false) +} + +/** The transformer geometry the KV-cache and forward-slab estimates need (from the GGUF header). */ +@ExperimentalMemoryApi +public data class ModelGeometry( + val layers: Int, + val heads: Int, + val kvHeads: Int, + /** Per-head key width (`attention.key_length`, = embedding / heads when absent). */ + val headDim: Int, + /** Per-head value width (`attention.value_length`, = [headDim] when absent). */ + val valueDim: Int = headDim, + val embeddingLength: Int, + /** Feed-forward width (`feed_forward_length`); 4 × embedding when absent. */ + val feedForwardLength: Int = 4 * embeddingLength, + val vocabSize: Int, + /** The context length the model was trained for, if the header says. */ + val trainedContextLength: Int? = null, +) + +/** How the KV cache is stored. */ +@ExperimentalMemoryApi +public enum class KvCacheMode(public val label: String) { + /** 2 bytes per element (bf16/f16). */ + BF16("bf16"), + /** TurboQuant 4-bit polar codes, block 128 (decision #11 default when the plan is tight). */ + TURBOQUANT_4("TurboQuant 4-bit"); + + /** Bytes for [elements] cache elements under this mode. */ + public fun bytes(elements: Long): Long = when (this) { + BF16 -> 2L * elements + TURBOQUANT_4 -> TensorEncoding.TurboQuantPolar(bitsPerElement = 4, blockSize = 128).physicalBytes(elements) ?: (elements / 2) + } +} + +/** What to plan for: the model (header only), the context length and the prefill chunk. */ +@ExperimentalMemoryApi +public data class PlanInput( + val modelName: String, + val architecture: String, + val weights: List, + val geometry: ModelGeometry?, + val ctx: Int, + val prefillChunk: Int = DEFAULT_PREFILL_CHUNK, + val kvMode: KvCacheMode = KvCacheMode.BF16, +) { + init { require(ctx > 0) { "ctx must be > 0" }; require(prefillChunk > 0) { "prefillChunk must be > 0" } } + + /** Names the checkpoint's name map could not translate — never dropped (M0-F4). */ + val unmappedWeights: List get() = weights.filter { it.id == null }.map { it.name } + + public companion object { + public const val DEFAULT_PREFILL_CHUNK: Int = 256 + } +} + +/** + * Memory budget the plan is checked against (decision #11): an explicit number of bytes, or + * `available − reserve` for a platform. + */ +@ExperimentalMemoryApi +public data class Budget(val bytes: Long, val description: String) { + public companion object { + /** Reserve the OS/app needs on Android and desktop JVMs (decision #11). */ + public const val RESERVE_ANDROID_JVM: Long = 700L * MiB + /** Reserve on Kotlin/Native targets (decision #11). */ + public const val RESERVE_NATIVE: Long = 300L * MiB + + public fun of(bytes: Long): Budget = Budget(bytes, "explicit") + public fun available(availableBytes: Long, reserve: Long = RESERVE_ANDROID_JVM): Budget = + Budget((availableBytes - reserve).coerceAtLeast(0), "available − reserve (${reserve / MiB} MB)") + } +} + +/** One line of the plan: what, how much, and whether it is resident for the whole session. */ +@ExperimentalMemoryApi +public data class PlanLine(val section: String, val detail: String, val bytes: Long, val resident: Boolean) + +/** A concrete way to make a plan fit, with the bytes it saves. */ +@ExperimentalMemoryApi +public data class Suggestion(val text: String, val savesBytes: Long) + +/** + * The memory plan of a model at a context length: weights (resident), KV cache, forward slab and + * heap headroom, totalled against a [Budget]. Pure arithmetic over [PlanInput]; nothing is + * allocated (PRD M0-F1..F3). The estimates are deliberately simple and documented in + * [MemoryPlans]; milestone M1's plan-vs-actual check calibrates them against real allocations. + */ +@ExperimentalMemoryApi +public data class MemoryPlan( + val input: PlanInput, + val weightsBytes: Long, + val kvBytes: Long, + /** KV bytes under the *other* mode, so the table can show both (bf16 vs TurboQuant). */ + val kvBytesAlternate: Long, + val forwardBytes: Long, + val headroomBytes: Long, + val budget: Budget?, +) { + val totalBytes: Long get() = weightsBytes + kvBytes + forwardBytes + headroomBytes + val residentBytes: Long get() = weightsBytes + kvBytes + + /** `true` when a budget is set and the total fits; `null` without a budget. */ + val fits: Boolean? get() = budget?.let { totalBytes <= it.bytes } + + val lines: List + get() = listOf( + PlanLine("weights", "Mapped, packed", weightsBytes, resident = true), + PlanLine("kv cache", input.kvMode.label + " @ ctx ${input.ctx}", kvBytes, resident = true), + PlanLine("forward", "prefill chunk ${input.prefillChunk}", forwardBytes, resident = false), + PlanLine("heap", "headroom", headroomBytes, resident = false), + ) + + /** At least two concrete suggestions with their savings when the plan does not fit (M0-F3). */ + public fun suggestions(): List { + val b = budget ?: return emptyList() + if (totalBytes <= b.bytes) return emptyList() + val out = ArrayList() + if (input.kvMode == KvCacheMode.BF16 && kvBytesAlternate < kvBytes) { + out += Suggestion("--kv turboquant", kvBytes - kvBytesAlternate) + } + val halfCtx = (input.ctx / 2).coerceAtLeast(1) + if (halfCtx < input.ctx) { + val half = MemoryPlans.plan(input.copy(ctx = halfCtx), budget) + out += Suggestion("--ctx $halfCtx", totalBytes - half.totalBytes) + } + val over = totalBytes - b.bytes + out += Suggestion("a smaller model: weights must shrink by ≥ ${MemoryPlans.formatBytes(over)} (e.g. a lower-bit quantization of the same model)", over) + return out + } + + /** The PRD §4.3 table. */ + public fun render(): String = buildString { + val g = input.geometry + append(input.modelName); append(" · "); append(input.architecture) + if (g != null) { append(" · "); append(g.layers); append(" layers") } + append(" · ctx "); append(input.ctx); append('\n') + for (l in lines) { + append(" "); append(l.section.padEnd(10)); append(l.detail.padEnd(26)); append(MemoryPlans.formatBytes(l.bytes).padStart(10)) + if (l.resident) append(" resident") + if (l.section == "kv cache") append(" (").append(MemoryPlans.formatBytes(kvBytesAlternate)).append(" with ").append(if (input.kvMode == KvCacheMode.BF16) KvCacheMode.TURBOQUANT_4.label else KvCacheMode.BF16.label).append(')') + append('\n') + } + append(" "); append("total".padEnd(36)); append(MemoryPlans.formatBytes(totalBytes).padStart(10)) + val b = budget + if (b != null) { + append(" of "); append(MemoryPlans.formatBytes(b.bytes)); append(if (fits == true) " ✔ fits" else " ✘ does not fit") + append('\n') + val s = suggestions() + if (s.isNotEmpty()) { + append(" suggestions: "); append(s.joinToString(" · ") { "${it.text} (−${MemoryPlans.formatBytes(it.savesBytes)})" }); append('\n') + } + } else append('\n') + val unmapped = input.unmappedWeights + if (unmapped.isNotEmpty()) { + append(" unmapped tensors (kept, not identified): "); append(unmapped.size); append(" — "); append(unmapped.take(5).joinToString(", ")); if (unmapped.size > 5) append(", …"); append('\n') + } + } +} + +/** The arithmetic behind [MemoryPlan]. */ +@ExperimentalMemoryApi +public object MemoryPlans { + + /** Fixed heap headroom the JVM/ART runtime needs besides tensors (decision #11 profile). */ + public const val HEAP_HEADROOM_BYTES: Long = 64L * MiB + + /** + * Build the plan. Estimates: + * - weights: sum of the packed byte sizes (they are touched every token, so counted resident); + * - kv cache: `layers × 2 × ctx × kvHeads × (headDim + valueDim)/2 × mode bytes`; + * - forward slab for a chunk of `T = min(prefillChunk, ctx)` tokens, FP32: + * `T × (4·emb + 3·ffn + heads·ctx) × 4 B` (residual stream, attention projections, gated FFN + * intermediates, attention scores over the context) plus one `vocab × 4 B` logits row; + * - heap headroom: [HEAP_HEADROOM_BYTES]. + */ + public fun plan(input: PlanInput, budget: Budget? = null): MemoryPlan { + val weights = input.weights.sumOf { it.bytes } + val g = input.geometry + val kvElements = if (g != null) kvElements(g, input.ctx) else 0L + val kv = input.kvMode.bytes(kvElements) + val kvAlt = (if (input.kvMode == KvCacheMode.BF16) KvCacheMode.TURBOQUANT_4 else KvCacheMode.BF16).bytes(kvElements) + val forward = if (g != null) forwardBytes(g, input.ctx, input.prefillChunk) else 0L + return MemoryPlan(input, weights, kv, kvAlt, forward, HEAP_HEADROOM_BYTES, budget) + } + + public fun kvElements(g: ModelGeometry, ctx: Int): Long = + g.layers.toLong() * ctx * g.kvHeads * (g.headDim + g.valueDim) + + public fun forwardBytes(g: ModelGeometry, ctx: Int, prefillChunk: Int): Long { + val t = minOf(prefillChunk, ctx).toLong() + val perToken = 4L * g.embeddingLength + 3L * g.feedForwardLength + g.heads.toLong() * ctx + return t * perToken * 4L + g.vocabSize.toLong() * 4L + } + + public fun formatBytes(bytes: Long): String = when { + bytes >= GiB -> formatOneDecimal(bytes.toDouble() / GiB) + " GB" + bytes >= MiB -> (bytes / MiB).toString() + " MB" + bytes >= KiB -> (bytes / KiB).toString() + " KB" + else -> "$bytes B" + } + + private fun formatOneDecimal(v: Double): String { + val tenths = kotlin.math.round(v * 10).toLong() + return "${tenths / 10}.${tenths % 10}" + } +} + +internal const val KiB: Long = 1024L +internal const val MiB: Long = 1024L * 1024L +internal const val GiB: Long = 1024L * 1024L * 1024L diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt new file mode 100644 index 000000000..4630a1748 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt @@ -0,0 +1,127 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ScopeKind +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.MemoryDomain +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** SKEEP-003 M0-A4: plan arithmetic, fit check and suggestions — from shapes and encodings only. */ +@OptIn(ExperimentalMemoryApi::class) +class MemoryPlanTest { + + /** Llama-3.2-1B-like geometry. */ + private val llama1b = ModelGeometry( + layers = 16, heads = 32, kvHeads = 8, headDim = 64, valueDim = 64, + embeddingLength = 2048, feedForwardLength = 8192, vocabSize = 128_256, trainedContextLength = 131_072, + ) + + private fun q4k(name: String, elements: Long): PlanTensor { + val f = Format(FP32, TensorEncoding.Q4_K) + return PlanTensor(name, TensorId.parse(name), f, elements, f.physicalBytes(elements)!!) + } + + private fun llamaWeights(): List = buildList { + add(q4k("model.embed_tokens.weight", 128_256L * 2048)) + for (n in 0 until 16) { + add(q4k("model.layers[$n].attn.q_proj.weight", 2048L * 2048)); add(q4k("model.layers[$n].attn.k_proj.weight", 512L * 2048)) + add(q4k("model.layers[$n].attn.v_proj.weight", 512L * 2048)); add(q4k("model.layers[$n].attn.o_proj.weight", 2048L * 2048)) + add(q4k("model.layers[$n].mlp.gate_proj.weight", 8192L * 2048)); add(q4k("model.layers[$n].mlp.up_proj.weight", 8192L * 2048)) + add(q4k("model.layers[$n].mlp.down_proj.weight", 2048L * 8192)) + } + } + + @Test + fun kvCacheArithmetic() { + // layers × 2 (K,V) × ctx × kvHeads × headDim × 2 B = 16 × 2 × 2048 × 8 × 64 × 2 = 64 MiB + assertEquals(16L * 2048 * 8 * 128, MemoryPlans.kvElements(llama1b, 2048)) + assertEquals(64L * MiB, KvCacheMode.BF16.bytes(MemoryPlans.kvElements(llama1b, 2048))) + val tq = KvCacheMode.TURBOQUANT_4.bytes(MemoryPlans.kvElements(llama1b, 2048)) + assertTrue(tq in (14L * MiB)..(20L * MiB), "TurboQuant 4-bit KV should be ~¼ of bf16, was ${tq / MiB} MB") + } + + @Test + fun forwardSlabScalesWithChunkAndContext() { + val decode = MemoryPlans.forwardBytes(llama1b, 2048, 1) + val prefill = MemoryPlans.forwardBytes(llama1b, 2048, 256) + assertTrue(prefill > decode) + // decode: 1 token × (4·2048 + 3·8192 + 32·2048) × 4 B + 128256 × 4 B + assertEquals((4L * 2048 + 3 * 8192 + 32 * 2048) * 4 + 128_256L * 4, decode) + assertEquals(MemoryPlans.forwardBytes(llama1b, 128, 256), MemoryPlans.forwardBytes(llama1b, 128, 1024)) // chunk capped at ctx + } + + @Test + fun planTotalsAndResidency() { + val input = PlanInput("Llama-3.2-1B-Instruct", "llama", llamaWeights(), llama1b, ctx = 2048) + val plan = MemoryPlans.plan(input, Budget.of(1300L * MiB)) + assertEquals(input.weights.sumOf { it.bytes }, plan.weightsBytes) + assertTrue(plan.weightsBytes in (600L * MiB)..(900L * MiB), "Q4_K weights ~0.7 GB, was ${plan.weightsBytes / MiB} MB") + assertEquals(64L * MiB, plan.kvBytes) + assertEquals(plan.weightsBytes + plan.kvBytes, plan.residentBytes) + assertEquals(plan.weightsBytes + plan.kvBytes + plan.forwardBytes + plan.headroomBytes, plan.totalBytes) + assertEquals(true, plan.fits) + assertTrue(plan.suggestions().isEmpty()) + assertEquals(listOf("weights", "kv cache", "forward", "heap"), plan.lines.map { it.section }) + assertTrue(plan.lines.first { it.section == "weights" }.resident) + assertFalse(plan.lines.first { it.section == "forward" }.resident) + // every weight's allocation is a mapped, model-lifetime, read-only spec + val a = input.weights.first().allocation + assertEquals(MemoryDomain.MMAP_FILE, a.domain); assertEquals(ScopeKind.MODEL, a.scope); assertFalse(a.mutable) + } + + @Test + fun doesNotFitGivesAtLeastTwoSuggestionsWithSavings() { + val input = PlanInput("Llama-3.2-1B-Instruct", "llama", llamaWeights(), llama1b, ctx = 2048) + val plan = MemoryPlans.plan(input, Budget.of(500L * MiB)) + assertEquals(false, plan.fits) + val s = plan.suggestions() + assertTrue(s.size >= 2, "need ≥ 2 suggestions, got $s") + assertTrue(s.any { it.text == "--kv turboquant" && it.savesBytes == plan.kvBytes - plan.kvBytesAlternate }) + assertTrue(s.any { it.text == "--ctx 1024" && it.savesBytes > 0 }) + assertTrue(s.all { it.savesBytes > 0 }) + val r = plan.render() + assertTrue(r.contains("✘ does not fit"), r) + assertTrue(r.contains("suggestions:"), r) + } + + @Test + fun turboQuantModeAndBudgetFromAvailableMemory() { + val input = PlanInput("m", "llama", llamaWeights(), llama1b, ctx = 2048, kvMode = KvCacheMode.TURBOQUANT_4) + val plan = MemoryPlans.plan(input, Budget.available(2L * GiB)) + assertEquals(2L * GiB - Budget.RESERVE_ANDROID_JVM, plan.budget!!.bytes) + assertTrue(plan.kvBytes < plan.kvBytesAlternate) + assertEquals(64L * MiB, plan.kvBytesAlternate) + assertTrue(plan.render().contains("TurboQuant 4-bit @ ctx 2048"), plan.render()) + } + + @Test + fun withoutGeometryOrBudget() { + val input = PlanInput("tiny", "test", listOf(q4k("w", 1024)), geometry = null, ctx = 512) + val plan = MemoryPlans.plan(input) + assertEquals(0L, plan.kvBytes); assertEquals(0L, plan.forwardBytes) + assertNull(plan.fits); assertTrue(plan.suggestions().isEmpty()) + assertEquals(input.weights.single().bytes + MemoryPlans.HEAP_HEADROOM_BYTES, plan.totalBytes) + assertTrue(plan.render().contains("tiny · test · ctx 512")) + } + + @Test + fun unmappedWeightsAreListedNotDropped() { + val f = Format(FP32, TensorEncoding.Q8_0) + val input = PlanInput("m", "test", listOf(PlanTensor("blk.0.ffn_gate_exps.weight", null, f, 32, 34)), null, 16) + assertEquals(listOf("blk.0.ffn_gate_exps.weight"), input.unmappedWeights) + assertTrue(MemoryPlans.plan(input).render().contains("unmapped tensors (kept, not identified): 1")) + } + + @Test + fun formatBytes() { + assertEquals("512 B", MemoryPlans.formatBytes(512)); assertEquals("4 KB", MemoryPlans.formatBytes(4096)) + assertEquals("68 MB", MemoryPlans.formatBytes(68L * MiB + 100)); assertEquals("1.3 GB", MemoryPlans.formatBytes(1300L * MiB)) + } +}