From b6a6824ea6dba4377d1c79c07f02b5227dbc27a9 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sat, 22 Aug 2026 22:50:55 +0200 Subject: [PATCH] =?UTF-8?q?feat(apps):=20skainet-plan=20CLI=20=E2=80=94=20?= =?UTF-8?q?memory=20plan=20of=20a=20GGUF=20from=20its=20header,=20budget?= =?UTF-8?q?=20fit,=20suggestions,=20tensor=20list=20by=20TensorId=20(SKEEP?= =?UTF-8?q?-003=20M0)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M0 sample (#1001, PRD §4.3, M0-F6 JVM): `skainet plan [--ctx N] [--budget 1.3G] [--kv bf16|turboquant] [--prefill-chunk N] [--list ] [--no-budget]` prints the PRD table — weights (resident), KV cache at ctx with the alternate mode, forward slab, heap headroom, total against the budget with ✔/✘ and ≥ 2 suggestions with savings — and lists tensors by TensorId ("TensorId · Format · n · bytes · ← gguf name"). Exit code 1 when the plan does not fit. Reads the GGUF header only. New Kotlin/JVM application module skainet-apps/skainet-plan (kotlinx-cli, pattern of skainet-tensor-tools), README with sample output, unit tests for size parsing, glob matching and the list renderer; stdout forced to UTF-8. The planner itself is commonMain (sk.ainet.lang.memory.plan, StreamingGGUFReader.planInput) and callable from Android/Native directly; a Kotlin/Native binary is a follow-up. Closes #1013 Co-Authored-By: Claude Fable 5 --- settings.gradle.kts | 1 + skainet-apps/skainet-plan/README.md | 34 ++++++++ skainet-apps/skainet-plan/build.gradle.kts | 26 ++++++ .../kotlin/sk/ainet/apps/plan/SkainetPlan.kt | 85 +++++++++++++++++++ .../sk/ainet/apps/plan/SkainetPlanTest.kt | 53 ++++++++++++ 5 files changed, 199 insertions(+) create mode 100644 skainet-apps/skainet-plan/README.md create mode 100644 skainet-apps/skainet-plan/build.gradle.kts create mode 100644 skainet-apps/skainet-plan/src/main/kotlin/sk/ainet/apps/plan/SkainetPlan.kt create mode 100644 skainet-apps/skainet-plan/src/test/kotlin/sk/ainet/apps/plan/SkainetPlanTest.kt diff --git a/settings.gradle.kts b/settings.gradle.kts index 4a67c2b21..9061397d9 100644 --- a/settings.gradle.kts +++ b/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") diff --git a/skainet-apps/skainet-plan/README.md b/skainet-apps/skainet-plan/README.md new file mode 100644 index 000000000..6f3f3e52c --- /dev/null +++ b/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 ` 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. diff --git a/skainet-apps/skainet-plan/build.gradle.kts b/skainet-apps/skainet-plan/build.gradle.kts new file mode 100644 index 000000000..d84af609e --- /dev/null +++ b/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() +} diff --git a/skainet-apps/skainet-plan/src/main/kotlin/sk/ainet/apps/plan/SkainetPlan.kt b/skainet-apps/skainet-plan/src/main/kotlin/sk/ainet/apps/plan/SkainetPlan.kt new file mode 100644 index 000000000..cf0d7f2a9 --- /dev/null +++ b/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 [--ctx N] [--budget 1.3G] [--kv bf16|turboquant] [--list ]` + * + * 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) { + 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') + } +} diff --git a/skainet-apps/skainet-plan/src/test/kotlin/sk/ainet/apps/plan/SkainetPlanTest.kt b/skainet-apps/skainet-plan/src/test/kotlin/sk/ainet/apps/plan/SkainetPlanTest.kt new file mode 100644 index 000000000..d6125f7b7 --- /dev/null +++ b/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 { 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) + } +}