feat(apps): skainet-plan CLI — memory plan of a GGUF from its header, budget fit, suggestions, tensor list by TensorId (SKEEP-003 M0, S0.8)

settings.gradle.kts

@@ -80,6 +80,7 @@ include("skainet-test:skainet-test-java")
 // ====== APPS
 include("skainet-apps:skainet-grayscale-cli")
 include("skainet-apps:skainet-tensor-tools")
+include("skainet-apps:skainet-plan")
 include("skainet-io:skainet-io-safetensors")
 include("skainet-io:skainet-io-iree-params")
 

skainet-apps/skainet-plan/README.md

@@ -0,0 +1,34 @@
+# skainet-plan — know before you load
+
+`skainet plan` answers *will this model fit on this device at this context length?* from a GGUF
+**header alone** (shapes, encodings and architecture metadata — no tensor bytes are read). It is
+the milestone-M0 sample of the SKaiNET memory architecture (SKEEP-003, #1001/#1013).
+
+```
+$ ./gradlew :skainet-apps:skainet-plan:run --args="Llama-3.2-1B-Instruct-Q4_K_M.gguf --ctx 2048 --budget 1.3G"
+Llama-3.2-1B-Instruct · llama · 16 layers · ctx 2048
+  weights   Mapped, packed                  762 MB   resident
+  kv cache  bf16 @ ctx 2048                  64 MB   resident   (17 MB with TurboQuant 4-bit)
+  forward   prefill chunk 256                47 MB
+  heap      headroom                         64 MB
+  total                                     938 MB   of 1.3 GB  ✔ fits
+
+$ … --budget 0.9G
+  total                                     938 MB   of 921 MB  ✘ does not fit
+  suggestions: --kv turboquant (−47 MB) · --ctx 1024 (−45 MB) · a smaller model: weights must shrink by ≥ 17 MB (…)
+
+$ … --list 'model.layers[3].*'
+  model.layers[3].attn.q_proj.weight   Float32/Q4_K      n=4194304       2 MB   ← blk.3.attn_q.weight
+  model.layers[3].attn.k_proj.weight   Float32/Q4_K      n=1048576     576 KB   ← blk.3.attn_k.weight
+  …
+```
+
+Options: `--ctx N` (default: the model's trained context length, else 2048) · `--budget 1.3G|900M|bytes`
+(default: the JVM's max heap minus the 700 MB reserve of the 2 GB device profile) · `--kv bf16|turboquant` ·
+`--prefill-chunk N` (default 256) · `--list <glob>` over `TensorId`s · `--no-budget`. Exit code 1 when the
+plan does not fit the budget.
+
+The planner itself (`sk.ainet.lang.memory.plan`, `StreamingGGUFReader.planInput`) is `commonMain`
+code usable from Android and Kotlin/Native directly; this module is the JVM command line. The
+numbers are estimates documented in `MemoryPlans` and calibrated by milestone M1's plan-vs-actual
+check; unmapped tensor names (unknown architectures) are listed, never dropped.

skainet-apps/skainet-plan/build.gradle.kts

@@ -0,0 +1,26 @@
+plugins {
+    alias(libs.plugins.jetbrainsKotlinJvm)
+    application
+}
+
+dependencies {
+    implementation(kotlin("stdlib"))
+    implementation(libs.kotlinx.cli)
+
+    implementation(project(":skainet-lang:skainet-lang-core"))
+    implementation(project(":skainet-io:skainet-io-core"))
+    implementation(project(":skainet-io:skainet-io-gguf"))
+
+    testImplementation(kotlin("test"))
+    testImplementation(project(":skainet-io:skainet-io-gguf"))
+}
+
+application {
+    mainClass.set("sk.ainet.apps.plan.SkainetPlanKt")
+    // the table uses · ✔ ✘ ← — print them as UTF-8 regardless of the terminal locale
+    applicationDefaultJvmArgs = listOf("-Dstdout.encoding=UTF-8", "-Dstderr.encoding=UTF-8", "-Dfile.encoding=UTF-8")
+}
+
+tasks.test {
+    useJUnitPlatform()
+}

skainet-apps/skainet-plan/src/main/kotlin/sk/ainet/apps/plan/SkainetPlan.kt

@@ -0,0 +1,85 @@
+@file:OptIn(ExperimentalMemoryApi::class)
+
+package sk.ainet.apps.plan
+
+import kotlinx.cli.ArgParser
+import kotlinx.cli.ArgType
+import kotlinx.cli.default
+import sk.ainet.io.JvmRandomAccessSource
+import sk.ainet.io.gguf.StreamingGGUFReader
+import sk.ainet.io.gguf.planInput
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.plan.Budget
+import sk.ainet.lang.memory.plan.KvCacheMode
+import sk.ainet.lang.memory.plan.MemoryPlan
+import sk.ainet.lang.memory.plan.MemoryPlans
+import sk.ainet.lang.memory.plan.PlanInput
+import java.io.File
+import kotlin.system.exitProcess
+
+/**
+ * `skainet plan <model.gguf> [--ctx N] [--budget 1.3G] [--kv bf16|turboquant] [--list <glob>]`
+ *
+ * Milestone M0 of SKEEP-003 ("know before you load"): prints the memory plan of a GGUF model at a
+ * context length — weights (resident), KV cache, forward slab, heap headroom — against a budget,
+ * with concrete suggestions when it does not fit, and lists tensors by `TensorId`. Reads the GGUF
+ * header only; no tensor bytes are touched.
+ */
+public fun main(args: Array<String>) {
+    val parser = ArgParser("skainet-plan")
+    val model by parser.argument(ArgType.String, fullName = "model", description = "Path to the GGUF file")
+    val ctx by parser.option(ArgType.Int, fullName = "ctx", description = "Context length to plan for (default: the model's trained context length, or 2048)")
+    val budget by parser.option(ArgType.String, fullName = "budget", description = "Memory budget, e.g. 1.3G, 900M, 1500000000; default: JVM max heap + direct memory estimate")
+    val kv by parser.option(ArgType.Choice(listOf("bf16", "turboquant"), { it }), fullName = "kv", description = "KV cache mode").default("bf16")
+    val prefill by parser.option(ArgType.Int, fullName = "prefill-chunk", description = "Prefill chunk size for the forward slab").default(PlanInput.DEFAULT_PREFILL_CHUNK)
+    val list by parser.option(ArgType.String, fullName = "list", description = "List tensors whose TensorId matches this glob, e.g. 'model.layers[3].*'")
+    val noBudget by parser.option(ArgType.Boolean, fullName = "no-budget", description = "Print the plan without a fit check").default(false)
+    parser.parse(args)
+
+    val file = File(model)
+    if (!file.isFile) { System.err.println("skainet plan: file not found: $model"); exitProcess(2) }
+
+    val kvMode = if (kv == "turboquant") KvCacheMode.TURBOQUANT_4 else KvCacheMode.BF16
+    val plan = JvmRandomAccessSource.open(file).use { src ->
+        val reader = StreamingGGUFReader.open(src)
+        val input = reader.planInput(ctx = ctx, prefillChunk = prefill, kvMode = kvMode)
+        val b = when {
+            noBudget -> null
+            budget != null -> Budget.of(parseBytes(budget!!))
+            else -> Budget.available(Runtime.getRuntime().maxMemory())
+        }
+        MemoryPlans.plan(input, b)
+    }
+    print(plan.render())
+    list?.let { glob -> print(renderList(plan, glob)) }
+    exitProcess(if (plan.fits == false) 1 else 0)
+}
+
+/** `1.3G`, `900M`, `64K`, `123456` → bytes (decimal suffixes are binary multiples, as the plan prints them). */
+internal fun parseBytes(text: String): Long {
+    val t = text.trim().uppercase().removeSuffix("B")
+    val mult = when (t.lastOrNull()) { 'G' -> 1L shl 30; 'M' -> 1L shl 20; 'K' -> 1L shl 10; else -> 1L }
+    val num = if (mult == 1L) t else t.dropLast(1)
+    val value = num.toDoubleOrNull() ?: throw IllegalArgumentException("Not a size: '$text' (use e.g. 1.3G, 900M)")
+    return (value * mult).toLong()
+}
+
+/** `model.layers[3].*` → regex; `*` matches anything, `?` one char, everything else literally. */
+internal fun globToRegex(glob: String): Regex =
+    Regex("^" + glob.split('*').joinToString(".*") { part -> part.split('?').joinToString(".") { Regex.escape(it) } } + "$")
+
+internal fun renderList(plan: MemoryPlan, glob: String): String = buildString {
+    val re = globToRegex(glob)
+    val rows = plan.input.weights.filter { w -> (w.id?.canonical ?: w.name).let { re.matches(it) } }
+    append('\n'); append("tensors matching '").append(glob).append("': ").append(rows.size).append('\n')
+    val idWidth = (rows.maxOfOrNull { (it.id?.canonical ?: "—").length } ?: 10).coerceAtMost(60)
+    for (w in rows) {
+        append("  ")
+        append((w.id?.canonical ?: "—").padEnd(idWidth)); append("  ")
+        append(w.format.toString().padEnd(18))
+        append("n=").append(w.elementCount.toString().padEnd(12))
+        append(MemoryPlans.formatBytes(w.bytes).padStart(8))
+        append("   ← ").append(w.name)
+        append('\n')
+    }
+}

skainet-apps/skainet-plan/src/test/kotlin/sk/ainet/apps/plan/SkainetPlanTest.kt

@@ -0,0 +1,53 @@
+@file:OptIn(ExperimentalMemoryApi::class)
+
+package sk.ainet.apps.plan
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.plan.Budget
+import sk.ainet.lang.memory.plan.MemoryPlans
+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.TensorId
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
+
+class SkainetPlanTest {
+
+    @Test
+    fun parsesSizes() {
+        assertEquals(1L shl 30, parseBytes("1G")); assertEquals((1.3 * (1L shl 30)).toLong(), parseBytes("1.3G"))
+        assertEquals(900L shl 20, parseBytes("900M")); assertEquals(900L shl 20, parseBytes("900MB"))
+        assertEquals(64L shl 10, parseBytes("64k")); assertEquals(123456L, parseBytes("123456"))
+        assertFailsWith<IllegalArgumentException> { parseBytes("lots") }
+    }
+
+    @Test
+    fun globMatching() {
+        val re = globToRegex("model.layers[3].*")
+        assertTrue(re.matches("model.layers[3].attn.q_proj.weight"))
+        assertTrue(!re.matches("model.layers[13].attn.q_proj.weight"))
+        assertTrue(globToRegex("*.weight").matches("model.norm.weight"))
+        assertTrue(globToRegex("model.layers[?].attn.*").matches("model.layers[7].attn.k_proj.bias"))
+    }
+
+    @Test
+    fun listRendersIdFormatShapeAndSourceName() {
+        val f = Format(FP32, TensorEncoding.Q4_K)
+        val w = PlanTensor("blk.3.attn_q.weight", TensorId.parse("model.layers[3].attn.q_proj.weight"), f, 2048L * 2048, f.physicalBytes(2048L * 2048)!!)
+        val other = PlanTensor("blk.4.attn_q.weight", TensorId.parse("model.layers[4].attn.q_proj.weight"), f, 2048L * 2048, f.physicalBytes(2048L * 2048)!!)
+        val g = ModelGeometry(16, 32, 8, 64, 64, 2048, 8192, 128_256)
+        val plan = MemoryPlans.plan(PlanInput("m", "llama", listOf(w, other), g, 2048), Budget.of(1300L shl 20))
+        val out = renderList(plan, "model.layers[3].*")
+        assertTrue(out.contains("tensors matching 'model.layers[3].*': 1"), out)
+        assertTrue(out.contains("model.layers[3].attn.q_proj.weight"), out)
+        assertTrue(out.contains("Float32/Q4_K"), out)
+        assertTrue(out.contains("← blk.3.attn_q.weight"), out)
+        assertTrue(!out.contains("layers[4]"), out)
+    }
+}