Skip to content

feat(memory): MemoryPlan from the GGUF header — weights, KV, forward slab, budget fit, suggestions (SKEEP-003 P1, S0.7) #1056

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 1 commit into from
Aug 23, 2026
Merged

feat(memory): MemoryPlan from the GGUF header — weights, KV, forward slab, budget fit, suggestions (SKEEP-003 P1, S0.7) #1056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/GgufMemoryPlan.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 (`<arch>.block_count`, `<arch>.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)
}
72 changes: 72 additions & 0 deletions skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/GgufMemoryPlanTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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())
}
}
}
179 changes: 179 additions & 0 deletions skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <init> (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 <init> (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 <init> (IIIIIIIILjava/lang/Integer;)V
public synthetic fun <init> (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 <init> (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 <init> (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 <init> (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 <init> (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 <init> (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 <init> (Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;)V
public synthetic fun <init> (Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
Expand Down
Loading
Loading