From 8a14dcd0e358ea690078419f46364e21ac088454 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 12:05:59 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20planner=20device=20profiles=20?= =?UTF-8?q?=E2=80=94=202=20GB=20mobile=20and=20desktop,=20with=20automatic?= =?UTF-8?q?=20KV=20quantization=20and=20a=20dequant=20limit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1039 (SKEEP-003 P5, S2.6, proposal §8 item 1, decision #11; M2-F6). Decision #11's numbers existed as prose. A 2 GB phone and a workstation do not want the same defaults, and the difference is not taste: on the phone the reserve is large relative to the budget, weights must be mapped, the KV cache has to shrink before the model does, and a dispatcher-inserted dequantization is a defect rather than a slow path. - `PlannerProfile`: reserve, prefill chunk, KV mode, the fraction of the budget above which KV auto-quantizes, the heap/off-heap threshold, the dequant warn limit and whether it is strict, and whether weights are mapped. `MOBILE_2GB` (700 MB reserved, prefill 256, off-heap ≥ 256 KB, KV → TurboQuant-4 over 80 %, dequant warns over 5 %, weights mapped), `DESKTOP` (today's behaviour: no automatic re-quantization, heap staging), `NATIVE` (the smaller 300 MB reserve), and `forDevice()` which picks mobile at or below 3 GB of RAM. - `profile.plan(input, availableBytes)` returns a `ProfiledPlan` carrying the decisions it made — the KV switch is a note, not a silent rewrite — and `requireFits(device)` refuses before anything is allocated, against the two pools of #1038 rather than one total. - `checkDequant(share)` turns #1035's `adapterShareOfBytesRead` into a verdict: fine, a warning naming the missing kernel, or a failure under `strict()`. - `skainet-plan --profile mobile|desktop|native` plans under the rules and prints them above the table, so a plan read months later says which rules produced it. On Android `AndroidGguf.profiledPlan(...)` defaults to `MOBILE_2GB`. CI and the gate gain a leg for plain-JVM modules: the skainet-plan CLI and the benchmark publisher have `test`, not `jvmTest`, so the repo-wide leg never reached them and their tests had never run in CI. PlannerProfileTest covers one rule per test — budget arithmetic, each default, the off-heap threshold at and under the boundary, the profile overriding the caller's prefill chunk, KV quantizing only when tight (same model, same ctx, two budgets), the desktop profile leaving the cache alone, refusal against a budget and against a device's two pools, dequant OK/warn/ strict, and profile selection by RAM. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- .github/workflows/build.yml | 7 +- scripts/pr-gate.sh | 5 + .../kotlin/sk/ainet/apps/plan/SkainetPlan.kt | 49 ++++- .../sk/ainet/apps/plan/SkainetPlanTest.kt | 27 +++ .../io/gguf/AndroidGgufLoadingHostTest.kt | 27 +++ .../sk/ainet/io/gguf/AndroidGgufLoading.kt | 27 +++ .../api/jvm/skainet-lang-core.api | 89 ++++++++ .../ainet/lang/memory/plan/PlannerProfile.kt | 197 ++++++++++++++++++ .../lang/memory/plan/PlannerProfileTest.kt | 194 +++++++++++++++++ 9 files changed, 611 insertions(+), 11 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlannerProfileTest.kt diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 19c836007..b03a5f1a0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -34,7 +34,12 @@ jobs: matrix: include: - name: jvm - tasks: jvmTest + # Plain-JVM modules (the skainet-plan CLI, the engine benchmark publisher) have a + # `test` task rather than `jvmTest`, so the repo-wide jvmTest never reaches them. + tasks: >- + jvmTest + :skainet-apps:skainet-plan:test + :skainet-backends:benchmarks:jvm-cpu-publish:test # verifyNpmPins guards the npm-* pins in gradle/libs.versions.toml against # lockfile drift; it belongs on the leg that already has the JS toolchain. - name: js-wasm diff --git a/scripts/pr-gate.sh b/scripts/pr-gate.sh index 49e0d123a..1ddd5d0bf 100755 --- a/scripts/pr-gate.sh +++ b/scripts/pr-gate.sh @@ -56,6 +56,11 @@ step "assemble" step "Java consumer API tests" "${GRADLE[@]}" :skainet-test:skainet-test-java:test +# Plain-JVM modules (CLI tools, the benchmark publisher) have `test`, not `jvmTest`, so the +# repo-wide jvmTest leg never reaches them — their tests would otherwise only ever run locally. +step "JVM tool tests (skainet-plan, engine benchmark publisher)" +"${GRADLE[@]}" :skainet-apps:skainet-plan:test :skainet-backends:benchmarks:jvm-cpu-publish:test + # The Android compilations have host-side (JVM) unit tests — the mmap weight path (#921) and the # Android loading facade (#1038). They compile against androidMain, so they are the only thing that # proves that code builds and runs on the Android variant; nothing else in the gate touches it. 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 index cf0d7f2a9..76a267660 100644 --- 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 @@ -14,16 +14,24 @@ 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 sk.ainet.lang.memory.plan.PlannerProfile +import sk.ainet.lang.memory.plan.ProfiledPlan import java.io.File import kotlin.system.exitProcess /** - * `skainet plan [--ctx N] [--budget 1.3G] [--kv bf16|turboquant] [--list ]` + * `skainet plan [--ctx N] [--budget 1.3G] [--kv bf16|turboquant] [--profile P] [--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. + * + * `--profile` (M2-F6, #1039) plans under a device profile instead of raw defaults: `mobile` is the + * 2 GB phone (700 MB reserved, weights mapped, KV quantized automatically once the plan passes 80 % + * of the budget), `desktop` keeps today's behaviour, `native` reserves the smaller Kotlin/Native + * amount. The profile and every decision it made are printed above the table, so a plan can be read + * back later without guessing which rules produced it. */ public fun main(args: Array) { val parser = ArgParser("skainet-plan") @@ -34,25 +42,46 @@ public fun main(args: Array) { 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) + val profileName by parser.option( + ArgType.Choice(listOf("none", "mobile", "desktop", "native"), { it }), + fullName = "profile", + description = "Device profile whose rules the plan follows (M2-F6): mobile = 2 GB phone, desktop, native", + ).default("none") 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 profile = profileFor(profileName) + val profiled: ProfiledPlan = 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()) + val available = budget?.let { parseBytes(it) } ?: Runtime.getRuntime().maxMemory() + if (profile != null) { + // A profile owns the reserve, so --budget names what the *device* has, not what the + // plan may use; without a profile the flag keeps its original meaning. + profile.plan(input, available) + } else { + val b = when { + noBudget -> null + budget != null -> Budget.of(available) + else -> Budget.available(available) + } + ProfiledPlan(PlannerProfile("none", reserveBytes = 0), MemoryPlans.plan(input, b), emptyList()) } - MemoryPlans.plan(input, b) } - print(plan.render()) - list?.let { glob -> print(renderList(plan, glob)) } - exitProcess(if (plan.fits == false) 1 else 0) + print(if (profile != null) profiled.render() else profiled.plan.render()) + list?.let { glob -> print(renderList(profiled.plan, glob)) } + exitProcess(if (profiled.plan.fits == false) 1 else 0) +} + +/** The profile behind a `--profile` value; `null` for "none" (the plan's own defaults). */ +internal fun profileFor(name: String): PlannerProfile? = when (name) { + "mobile" -> PlannerProfile.MOBILE_2GB + "desktop" -> PlannerProfile.DESKTOP + "native" -> PlannerProfile.NATIVE + else -> null } /** `1.3G`, `900M`, `64K`, `123456` → bytes (decimal suffixes are binary multiples, as the plan prints them). */ 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 index d6125f7b7..1dd17b5c6 100644 --- 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 @@ -9,6 +9,7 @@ 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.memory.plan.PlannerProfile import sk.ainet.lang.tensor.TensorId import sk.ainet.lang.tensor.storage.TensorEncoding import sk.ainet.lang.types.FP32 @@ -50,4 +51,30 @@ class SkainetPlanTest { assertTrue(out.contains("← blk.3.attn_q.weight"), out) assertTrue(!out.contains("layers[4]"), out) } + + @Test + fun profileFlagSelectsTheDeviceRules() { + assertEquals(PlannerProfile.MOBILE_2GB, profileFor("mobile")) + assertEquals(PlannerProfile.DESKTOP, profileFor("desktop")) + assertEquals(PlannerProfile.NATIVE, profileFor("native")) + assertEquals(null, profileFor("none"), "no profile means the plan's own defaults, as before") + } + + @Test + fun theProfiledPlanPrintsItsRulesAboveTheTable() { + // #1039: a plan printed months later has to say which rules produced it. + val f = Format(FP32, TensorEncoding.Q4_K) + val w = PlanTensor("blk.0.attn_q.weight", null, f, 600L * (1 shl 20) / 144 * 256, 600L shl 20) + val g = ModelGeometry(16, 32, 8, 64, 64, 2048, 5632, 32_000) + val input = PlanInput("m", "llama", listOf(w), g, ctx = 8192) + + val mobile = PlannerProfile.MOBILE_2GB.plan(input, availableBytes = 2048L shl 20).render() + assertTrue(mobile.startsWith("profile mobile-2gb"), mobile) + assertTrue(mobile.contains("prefill 256"), mobile) + assertTrue(mobile.contains("note: KV cache switched"), "a tight mobile plan quantizes the cache:\n$mobile") + + val desktop = PlannerProfile.DESKTOP.plan(input, availableBytes = 2048L shl 20).render() + assertTrue(desktop.startsWith("profile desktop"), desktop) + assertTrue(!desktop.contains("switched"), "the desktop profile keeps today's behaviour:\n$desktop") + } } diff --git a/skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/AndroidGgufLoadingHostTest.kt b/skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/AndroidGgufLoadingHostTest.kt index bf40b3d1b..1e64d5319 100644 --- a/skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/AndroidGgufLoadingHostTest.kt +++ b/skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/AndroidGgufLoadingHostTest.kt @@ -6,6 +6,7 @@ import sk.ainet.io.model.QuantPolicy import sk.ainet.io.model.StagingPolicy import sk.ainet.lang.memory.ExperimentalMemoryApi import sk.ainet.lang.memory.plan.DeviceMemory +import sk.ainet.lang.memory.plan.PlannerProfile import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.FloatArrayTensorData import sk.ainet.lang.tensor.data.MmapFloatTensorData @@ -17,6 +18,7 @@ import java.nio.ByteOrder import kotlin.test.Test import kotlin.test.assertContentEquals import kotlin.test.assertEquals +import kotlin.test.assertFailsWith import kotlin.test.assertFalse import kotlin.test.assertTrue @@ -123,6 +125,31 @@ class AndroidGgufLoadingHostTest { } } + @Test + fun `the android plan follows the mobile profile by default`() { + val f = model() + try { + val phone = DeviceMemory( + totalRamBytes = 2048 * mb, availableRamBytes = 1600 * mb, + heapMaxBytes = 512 * mb, heapUsedBytes = 40 * mb, lowMemoryThresholdBytes = 180 * mb, + ) + val profiled = AndroidGguf.profiledPlan(f.absolutePath, ctx = 512, device = phone) + assertEquals(PlannerProfile.MOBILE_2GB, profiled.profile, "Android plans as a 2 GB phone (#1039)") + assertEquals(256, profiled.plan.input.prefillChunk) + assertTrue(profiled.render().contains("profile mobile-2gb"), profiled.render()) + profiled.requireFits(phone) + + // and it refuses, before allocating anything, when the device cannot take it + val tiny = phone.copy(availableRamBytes = 780 * mb, heapMaxBytes = 8 * mb, heapUsedBytes = 7 * mb) + val refusal = assertFailsWith { + AndroidGguf.profiledPlan(f.absolutePath, ctx = 8192, device = tiny).requireFits(tiny) + } + assertTrue(refusal.message!!.contains("mobile-2gb"), refusal.message!!) + } finally { + f.delete() + } + } + @Test fun `an unmapped load is charged for its weights, a mapped one is not`() { val f = model() diff --git a/skainet-io/skainet-io-gguf/src/androidMain/kotlin/sk/ainet/io/gguf/AndroidGgufLoading.kt b/skainet-io/skainet-io-gguf/src/androidMain/kotlin/sk/ainet/io/gguf/AndroidGgufLoading.kt index e6d546448..dcf2c9ef3 100644 --- a/skainet-io/skainet-io-gguf/src/androidMain/kotlin/sk/ainet/io/gguf/AndroidGgufLoading.kt +++ b/skainet-io/skainet-io-gguf/src/androidMain/kotlin/sk/ainet/io/gguf/AndroidGgufLoading.kt @@ -12,6 +12,8 @@ import sk.ainet.lang.memory.plan.DeviceFit import sk.ainet.lang.memory.plan.DeviceMemory import sk.ainet.lang.memory.plan.MemoryPlan import sk.ainet.lang.memory.plan.MemoryPlans +import sk.ainet.lang.memory.plan.PlannerProfile +import sk.ainet.lang.memory.plan.ProfiledPlan import sk.ainet.lang.memory.plan.fitOn /** @@ -94,4 +96,29 @@ public object AndroidGguf { /** [fits] against an explicit [DeviceMemory] — the form a test or a simulation uses. */ public fun fits(device: DeviceMemory, filePath: String, ctx: Int, weightsMapped: Boolean = true): DeviceFit = plan(filePath, ctx).fitOn(device, weightsMapped) + + /** + * The plan under a device profile (#1039): [PlannerProfile.MOBILE_2GB] by default on Android — + * 700 MB reserved, weights mapped, the KV cache quantized automatically once the plan passes + * 80 % of the budget — with every decision it made recorded in the result. + * + * `profiledPlan(...).requireFits(device)` is the refusal that happens *before* a byte is + * allocated (M2-F6). + */ + public fun profiledPlan( + filePath: String, + ctx: Int, + device: DeviceMemory, + profile: PlannerProfile = PlannerProfile.MOBILE_2GB, + ): ProfiledPlan = openSource(filePath).use { source -> + profile.plan(StreamingGGUFReader.open(source).planInput(ctx), device.availableRamBytes) + } + + /** [profiledPlan] reading this device's memory itself. */ + public fun profiledPlan( + context: Context, + filePath: String, + ctx: Int, + profile: PlannerProfile = PlannerProfile.MOBILE_2GB, + ): ProfiledPlan = profiledPlan(filePath, ctx, deviceMemory(context), profile) } 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 6ea9024f2..7297a5242 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 @@ -1295,6 +1295,33 @@ public final class sk/ainet/lang/memory/plan/Budget$Companion { public final fun of (J)Lsk/ainet/lang/memory/plan/Budget; } +public final class sk/ainet/lang/memory/plan/DequantSeverity : java/lang/Enum { + public static final field ERROR Lsk/ainet/lang/memory/plan/DequantSeverity; + public static final field OK Lsk/ainet/lang/memory/plan/DequantSeverity; + public static final field WARN Lsk/ainet/lang/memory/plan/DequantSeverity; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/memory/plan/DequantSeverity; + public static fun values ()[Lsk/ainet/lang/memory/plan/DequantSeverity; +} + +public final class sk/ainet/lang/memory/plan/DequantVerdict { + public fun (Lsk/ainet/lang/memory/plan/PlannerProfile;DLsk/ainet/lang/memory/plan/DequantSeverity;Ljava/lang/String;)V + public final fun component1 ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun component2 ()D + public final fun component3 ()Lsk/ainet/lang/memory/plan/DequantSeverity; + public final fun component4 ()Ljava/lang/String; + public final fun copy (Lsk/ainet/lang/memory/plan/PlannerProfile;DLsk/ainet/lang/memory/plan/DequantSeverity;Ljava/lang/String;)Lsk/ainet/lang/memory/plan/DequantVerdict; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/DequantVerdict;Lsk/ainet/lang/memory/plan/PlannerProfile;DLsk/ainet/lang/memory/plan/DequantSeverity;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/DequantVerdict; + public fun equals (Ljava/lang/Object;)Z + public final fun getMessage ()Ljava/lang/String; + public final fun getProfile ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun getSeverity ()Lsk/ainet/lang/memory/plan/DequantSeverity; + public final fun getShare ()D + public fun hashCode ()I + public final fun requireAcceptable ()V + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/memory/plan/DeviceFit { public fun (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/DeviceMemory;ZLsk/ainet/lang/memory/plan/PoolFit;Lsk/ainet/lang/memory/plan/PoolFit;Ljava/util/List;)V public final fun component1 ()Lsk/ainet/lang/memory/plan/MemoryPlan; @@ -1547,6 +1574,49 @@ public final class sk/ainet/lang/memory/plan/PlanVsActualLine { public final fun withinTolerance (D)Z } +public final class sk/ainet/lang/memory/plan/PlannerProfile { + public static final field Companion Lsk/ainet/lang/memory/plan/PlannerProfile$Companion; + public static final field MOBILE_RAM_CEILING J + public static final field OFF_HEAP_THRESHOLD J + public fun (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZ)V + public synthetic fun (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun budget (J)Lsk/ainet/lang/memory/plan/Budget; + public final fun checkDequant (Ljava/lang/Double;)Lsk/ainet/lang/memory/plan/DequantVerdict; + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()J + public final fun component3 ()I + public final fun component4 ()Lsk/ainet/lang/memory/plan/KvCacheMode; + public final fun component5 ()D + public final fun component6 ()J + public final fun component7 ()D + public final fun component8 ()Z + public final fun component9 ()Z + public final fun copy (Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZ)Lsk/ainet/lang/memory/plan/PlannerProfile; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlannerProfile;Ljava/lang/String;JILsk/ainet/lang/memory/plan/KvCacheMode;DJDZZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun domainFor (J)Lsk/ainet/lang/tensor/storage/MemoryDomain; + public fun equals (Ljava/lang/Object;)Z + public final fun getDequantWarnFraction ()D + public final fun getKvAutoQuantizeAbove ()D + public final fun getKvMode ()Lsk/ainet/lang/memory/plan/KvCacheMode; + public final fun getName ()Ljava/lang/String; + public final fun getOffHeapThresholdBytes ()J + public final fun getPrefillChunk ()I + public final fun getReserveBytes ()J + public final fun getStrict ()Z + public final fun getWeightsMapped ()Z + public fun hashCode ()I + public final fun plan (Lsk/ainet/lang/memory/plan/PlanInput;J)Lsk/ainet/lang/memory/plan/ProfiledPlan; + public final fun strict ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public fun toString ()Ljava/lang/String; +} + +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 getMOBILE_2GB ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun getNATIVE ()Lsk/ainet/lang/memory/plan/PlannerProfile; +} + public final class sk/ainet/lang/memory/plan/PoolFit { public fun (Ljava/lang/String;JJ)V public final fun component1 ()Ljava/lang/String; @@ -1564,6 +1634,25 @@ public final class sk/ainet/lang/memory/plan/PoolFit { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/plan/ProfiledPlan { + public fun (Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/MemoryPlan;Ljava/util/List;)V + public final fun component1 ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public final fun component2 ()Lsk/ainet/lang/memory/plan/MemoryPlan; + public final fun component3 ()Ljava/util/List; + public final fun copy (Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/MemoryPlan;Ljava/util/List;)Lsk/ainet/lang/memory/plan/ProfiledPlan; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/ProfiledPlan;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/MemoryPlan;Ljava/util/List;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/ProfiledPlan; + public fun equals (Ljava/lang/Object;)Z + public final fun getFits ()Ljava/lang/Boolean; + public final fun getNotes ()Ljava/util/List; + public final fun getPlan ()Lsk/ainet/lang/memory/plan/MemoryPlan; + public final fun getProfile ()Lsk/ainet/lang/memory/plan/PlannerProfile; + public fun hashCode ()I + public final fun render ()Ljava/lang/String; + public final fun requireFits (Lsk/ainet/lang/memory/plan/DeviceMemory;)V + public static synthetic fun requireFits$default (Lsk/ainet/lang/memory/plan/ProfiledPlan;Lsk/ainet/lang/memory/plan/DeviceMemory;ILjava/lang/Object;)V + 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; 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 new file mode 100644 index 000000000..fff38fbdc --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlannerProfile.kt @@ -0,0 +1,197 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.storage.MemoryDomain + +/** + * The rules a planner applies for a class of device (SKEEP-003 §8 item 1, decision #11; M2-F6). + * + * A 2 GB phone and a workstation do not want the same defaults, and the difference is not a matter + * of taste: on the phone the reserve is large relative to the budget, weights *must* be mapped, the + * KV cache has to shrink before the model does, and a dispatcher-inserted dequantization is a bug + * rather than a slow path. A profile makes those rules explicit and testable instead of leaving + * them as folklore in whoever configures the loader. + * + * @property reserveBytes what the OS and the rest of the app need; the budget is `available − this` + * @property prefillChunk tokens per prefill step — the width the forward slab is pre-sized for + * @property kvMode the KV format a comfortable plan uses + * @property kvAutoQuantizeAbove switch KV to [KvCacheMode.TURBOQUANT_4] when the plan needs more + * than this fraction of the budget (`1.0` disables the rule) + * @property offHeapThresholdBytes allocations of at least this many bytes belong off the heap + * @property dequantWarnFraction dispatcher-inserted dequantization above this share of the bytes a + * decode step reads is worth a warning — it means a kernel is missing for the format on disk + * @property strict turn those warnings into failures + * @property weightsMapped weights are expected to load through `StagingPolicy.MAPPED` + */ +@ExperimentalMemoryApi +public data class PlannerProfile( + val name: String, + val reserveBytes: Long, + val prefillChunk: Int = PlanInput.DEFAULT_PREFILL_CHUNK, + val kvMode: KvCacheMode = KvCacheMode.BF16, + val kvAutoQuantizeAbove: Double = 1.0, + val offHeapThresholdBytes: Long = OFF_HEAP_THRESHOLD, + val dequantWarnFraction: Double = 0.05, + val strict: Boolean = false, + val weightsMapped: Boolean = false, +) { + init { + require(reserveBytes >= 0) { "reserveBytes must be >= 0" } + require(prefillChunk > 0) { "prefillChunk must be > 0" } + require(kvAutoQuantizeAbove > 0.0) { "kvAutoQuantizeAbove must be > 0" } + require(offHeapThresholdBytes > 0) { "offHeapThresholdBytes must be > 0" } + } + + /** The budget this profile derives from [availableBytes]. */ + public fun budget(availableBytes: Long): Budget = + Budget((availableBytes - reserveBytes).coerceAtLeast(0L), "$name: available − ${MemoryPlans.formatBytes(reserveBytes)}") + + /** Where an allocation of [bytes] belongs — the heap/off-heap threshold (decision #11). */ + public fun domainFor(bytes: Long): MemoryDomain = + if (bytes >= offHeapThresholdBytes) MemoryDomain.HOST_OFFHEAP else MemoryDomain.HOST_HEAP + + /** + * Plan [input] under this profile against [availableBytes]. + * + * The profile's prefill chunk and KV mode replace whatever the input carried, and if the + * resulting plan needs more than [kvAutoQuantizeAbove] of the budget, the KV cache is + * re-planned as TurboQuant-4 *before* anything else is suggested — quantizing the cache is the + * cheapest thing to give up, and doing it automatically is the difference between a model that + * loads and a error message on a phone. + */ + public fun plan(input: PlanInput, availableBytes: Long): ProfiledPlan { + val budget = budget(availableBytes) + val base = MemoryPlans.plan(input.copy(prefillChunk = prefillChunk, kvMode = kvMode), budget) + val notes = ArrayList() + var plan = base + val share = if (budget.bytes > 0) base.totalBytes.toDouble() / budget.bytes else Double.MAX_VALUE + if (kvMode != KvCacheMode.TURBOQUANT_4 && share > kvAutoQuantizeAbove) { + val quantized = MemoryPlans.plan(base.input.copy(kvMode = KvCacheMode.TURBOQUANT_4), budget) + if (quantized.totalBytes < base.totalBytes) { + plan = quantized + notes += "KV cache switched to ${KvCacheMode.TURBOQUANT_4.label}: the plan needed " + + "${percent(share)} of the budget (over ${percent(kvAutoQuantizeAbove)}), saving " + + MemoryPlans.formatBytes(base.kvBytes - quantized.kvBytes) + } + } + if (weightsMapped) notes += "weights are counted resident and mapped — off the managed heap" + return ProfiledPlan(this, plan, notes) + } + + /** + * The verdict on how much the dispatcher had to dequantize during decode, given the share of + * bytes read that went through an adapter (`GenerationMetrics.adapterShareOfBytesRead`). + * + * A packed weight that has to be widened before every matmul means the kernel for its format is + * missing — the cost is real memory traffic, not a rounding error, which is why this profile + * can be run [strict] in CI and lenient on a desktop. + */ + public fun checkDequant(adapterShareOfBytesRead: Double?): DequantVerdict { + val share = adapterShareOfBytesRead ?: return DequantVerdict(this, 0.0, DequantSeverity.OK, "no bytes read") + if (share <= dequantWarnFraction) { + return DequantVerdict(this, share, DequantSeverity.OK, "adapters moved ${percent(share)} of the bytes read") + } + val message = "adapters moved ${percent(share)} of the bytes read, over the ${percent(dequantWarnFraction)} " + + "$name limit — a kernel for the on-disk format is missing and every step pays for it" + return DequantVerdict(this, share, if (strict) DequantSeverity.ERROR else DequantSeverity.WARN, message) + } + + /** This profile with [strict] on — what a CI run or an acceptance test uses. */ + public fun strict(): PlannerProfile = copy(name = "$name (strict)", strict = true) + + public companion object { + /** Off-heap threshold from decision #11: below this, a heap array is cheaper than a mapping. */ + public const val OFF_HEAP_THRESHOLD: Long = 256L * 1024 + + /** + * A 2 GB-class phone: 700 MB reserved for the OS and the app, weights mapped, KV + * automatically quantized once the plan passes 80 % of the budget, and dequantization + * treated as the defect it is. The default on Android. + */ + public val MOBILE_2GB: PlannerProfile = PlannerProfile( + name = "mobile-2gb", + reserveBytes = Budget.RESERVE_ANDROID_JVM, + prefillChunk = PlanInput.DEFAULT_PREFILL_CHUNK, + kvMode = KvCacheMode.BF16, + 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", + reserveBytes = Budget.RESERVE_ANDROID_JVM, + ) + + /** Kotlin/Native hosts, which reserve less than a JVM (decision #11). */ + public val NATIVE: PlannerProfile = PlannerProfile( + name = "native", + reserveBytes = Budget.RESERVE_NATIVE, + ) + + /** The profile a device's own numbers call for: mobile below [MOBILE_RAM_CEILING] of RAM. */ + public fun forDevice(device: DeviceMemory): PlannerProfile = + if (device.totalRamBytes <= MOBILE_RAM_CEILING) MOBILE_2GB else DESKTOP + + /** Devices with at most this much RAM get [MOBILE_2GB]. */ + public const val MOBILE_RAM_CEILING: Long = 3L * 1024 * MiB + } +} + +private fun percent(fraction: Double): String { + val scaled = (fraction * 1000).toLong() + return "${scaled / 10}.${scaled % 10}%" +} + +/** How bad a dequantization share is under a profile. */ +@ExperimentalMemoryApi +public enum class DequantSeverity { OK, WARN, ERROR } + +/** The verdict of [PlannerProfile.checkDequant]. */ +@ExperimentalMemoryApi +public data class DequantVerdict( + val profile: PlannerProfile, + val share: Double, + val severity: DequantSeverity, + val message: String, +) { + /** Throw when this profile is strict and the share is over its limit. */ + public fun requireAcceptable() { + if (severity == DequantSeverity.ERROR) throw IllegalStateException(message) + } +} + +/** A [MemoryPlan] made under a [PlannerProfile], with what the profile decided along the way. */ +@ExperimentalMemoryApi +public data class ProfiledPlan( + val profile: PlannerProfile, + val plan: MemoryPlan, + val notes: List, +) { + val fits: Boolean? get() = plan.fits + + /** + * Refuse before anything is allocated (M2-F6): throws when the plan does not fit its budget, or + * when [device] is given and either of its pools is short. + */ + public fun requireFits(device: DeviceMemory? = null) { + if (device != null) { + val fit = plan.fitOn(device, profile.weightsMapped) + if (!fit.fits) throw IllegalStateException("does not fit this device under profile '${profile.name}':\n" + fit.render()) + return + } + if (plan.fits == false) throw IllegalStateException("does not fit the budget under profile '${profile.name}':\n" + plan.render()) + } + + public fun render(): String = buildString { + append("profile ").append(profile.name) + append(" · reserve ").append(MemoryPlans.formatBytes(profile.reserveBytes)) + append(" · prefill ").append(profile.prefillChunk) + append(" · off-heap ≥ ").append(MemoryPlans.formatBytes(profile.offHeapThresholdBytes)) + if (profile.kvAutoQuantizeAbove < 1.0) append(" · KV auto-quantize over ").append(percent(profile.kvAutoQuantizeAbove)) + if (profile.strict) append(" · strict") + append('\n') + for (n in notes) append(" note: ").append(n).append('\n') + append(plan.render()) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlannerProfileTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlannerProfileTest.kt new file mode 100644 index 000000000..09d358864 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlannerProfileTest.kt @@ -0,0 +1,194 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +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.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * #1039 (M2-F6, decision #11): the planner's device profiles — one test per rule, because each of + * them is a number someone will otherwise "improve" by feel. + */ +@OptIn(ExperimentalMemoryApi::class) +class PlannerProfileTest { + + private val mb = 1024L * 1024L + + private fun input(ctx: Int = 2048, weightsMb: Long = 600, prefillChunk: Int = 64): PlanInput { + val f = Format(FP32, TensorEncoding.Q4_K) + val elements = weightsMb * mb / 144 * 256 + return PlanInput( + modelName = "llama-1b", + architecture = "llama", + weights = listOf(PlanTensor("model.weight", null, f, elements, weightsMb * mb)), + geometry = ModelGeometry( + layers = 16, heads = 32, kvHeads = 8, headDim = 64, + embeddingLength = 2048, feedForwardLength = 5632, vocabSize = 32000, + ), + ctx = ctx, + prefillChunk = prefillChunk, + kvMode = KvCacheMode.FP32, + ) + } + + // --- budget ------------------------------------------------------------------------------ + + @Test + fun theBudgetIsAvailableMinusTheProfilesReserve() { + val available = 2048 * mb + assertEquals( + available - Budget.RESERVE_ANDROID_JVM, + PlannerProfile.MOBILE_2GB.budget(available).bytes, + "mobile reserves 700 MB for the OS and the app (decision #11)", + ) + assertEquals( + available - Budget.RESERVE_NATIVE, + PlannerProfile.NATIVE.budget(available).bytes, + "Kotlin/Native reserves 300 MB", + ) + assertEquals(0L, PlannerProfile.MOBILE_2GB.budget(100 * mb).bytes, "never negative") + assertTrue(PlannerProfile.MOBILE_2GB.budget(available).description.contains("mobile-2gb")) + } + + // --- the profile's own defaults ----------------------------------------------------------- + + @Test + fun theProfilesDefaultsAreTheOnesDecisionElevenNames() { + val m = PlannerProfile.MOBILE_2GB + assertEquals(256, m.prefillChunk, "prefill chunked at 256 tokens") + assertEquals(256L * 1024, m.offHeapThresholdBytes, "heap/off-heap threshold is 256 KB") + assertEquals(0.80, m.kvAutoQuantizeAbove, "KV auto-quantizes over 80 % of the budget") + assertEquals(0.05, m.dequantWarnFraction, "dispatcher dequant warns over 5 %") + assertTrue(m.weightsMapped, "on a phone the weights are mapped") + assertFalse(m.strict) + + val d = PlannerProfile.DESKTOP + assertEquals(1.0, d.kvAutoQuantizeAbove, "a desktop never silently re-quantizes the cache") + assertFalse(d.weightsMapped, "desktop staging is the heap — today's behaviour, unchanged") + } + + @Test + fun theOffHeapThresholdDecidesWhereAnAllocationGoes() { + val p = PlannerProfile.MOBILE_2GB + assertEquals(MemoryDomain.HOST_HEAP, p.domainFor(255L * 1024)) + assertEquals(MemoryDomain.HOST_OFFHEAP, p.domainFor(256L * 1024), "at the threshold, not over it") + assertEquals(MemoryDomain.HOST_OFFHEAP, p.domainFor(64 * mb)) + } + + // --- the forward slab and the KV cache ----------------------------------------------------- + + @Test + fun theProfileImposesItsPrefillChunkOnTheInput() { + val profiled = PlannerProfile.MOBILE_2GB.plan(input(prefillChunk = 8), availableBytes = 4096 * mb) + assertEquals(256, profiled.plan.input.prefillChunk, "the profile's slab width wins over the caller's") + } + + @Test + fun kvIsQuantizedAutomaticallyOnlyWhenThePlanIsTight() { + // the same model at the same context length, planned twice: only the budget differs. + val model = input(ctx = 8192) + + // comfortable: 8 GB available — the cache stays as planned + val roomy = PlannerProfile.MOBILE_2GB.plan(model, availableBytes = 8192 * mb) + assertEquals(KvCacheMode.BF16, roomy.plan.input.kvMode, roomy.render()) + assertTrue(roomy.notes.none { it.contains("switched") }, roomy.notes.toString()) + + // tight: a 2 GB device — the cache quantizes before anything else gives + val tight = PlannerProfile.MOBILE_2GB.plan(model, availableBytes = 2048 * mb) + assertEquals(KvCacheMode.TURBOQUANT_4, tight.plan.input.kvMode, tight.render()) + assertTrue(tight.notes.any { it.contains("KV cache switched") }, tight.notes.toString()) + assertTrue( + tight.plan.kvBytes < roomy.plan.kvBytes, + "quantizing must shrink the cache: ${tight.plan.kvBytes} vs ${roomy.plan.kvBytes}", + ) + assertTrue(tight.plan.totalBytes < roomy.plan.totalBytes, "and the plan with it") + } + + @Test + fun theDesktopProfileLeavesTheCacheAlone() { + val tight = PlannerProfile.DESKTOP.plan(input(ctx = 8192), availableBytes = 2048 * mb) + assertEquals(KvCacheMode.BF16, tight.plan.input.kvMode, "no automatic re-quantization off-device") + assertTrue(tight.notes.isEmpty()) + } + + // --- the fit check ------------------------------------------------------------------------ + + @Test + fun theFitCheckRefusesBeforeAnythingIsAllocated() { + val profiled = PlannerProfile.MOBILE_2GB.plan(input(weightsMb = 1600), availableBytes = 2048 * mb) + assertEquals(false, profiled.fits) + val failure = assertFailsWith { profiled.requireFits() } + assertTrue(failure.message!!.contains("mobile-2gb"), failure.message!!) + assertTrue(profiled.plan.suggestions().isNotEmpty(), "and it says what to do instead") + } + + @Test + fun theFitCheckCanRefuseAgainstARealDevicesTwoPools() { + val phone = DeviceMemory( + totalRamBytes = 2048 * mb, availableRamBytes = 900 * mb, + heapMaxBytes = 512 * mb, heapUsedBytes = 40 * mb, lowMemoryThresholdBytes = 180 * mb, + ) + val mobile = PlannerProfile.MOBILE_2GB.plan(input(weightsMb = 600, ctx = 512), availableBytes = 900 * mb) + mobile.requireFits(phone) // mapped weights: the heap only carries KV + forward + headroom + + // the same plan with heap staging is charged for the weights and cannot fit + val onHeap = PlannerProfile.DESKTOP.plan(input(weightsMb = 600, ctx = 512), availableBytes = 900 * mb) + val failure = assertFailsWith { onHeap.requireFits(phone) } + assertTrue(failure.message!!.contains("managed heap"), failure.message!!) + } + + // --- dispatcher-inserted dequantization ---------------------------------------------------- + + @Test + fun dequantizationUnderTheLimitIsFine() { + val verdict = PlannerProfile.MOBILE_2GB.checkDequant(0.04) + assertEquals(DequantSeverity.OK, verdict.severity, verdict.message) + verdict.requireAcceptable() + assertEquals(DequantSeverity.OK, PlannerProfile.MOBILE_2GB.checkDequant(null).severity, "nothing read, nothing to judge") + } + + @Test + fun dequantizationOverTheLimitWarnsAndFailsUnderStrict() { + val lenient = PlannerProfile.MOBILE_2GB.checkDequant(0.31) + assertEquals(DequantSeverity.WARN, lenient.severity) + assertTrue(lenient.message.contains("31.0%"), lenient.message) + assertTrue(lenient.message.contains("kernel for the on-disk format is missing"), lenient.message) + lenient.requireAcceptable() // a warning does not stop a desktop run + + val strict = PlannerProfile.MOBILE_2GB.strict().checkDequant(0.31) + assertEquals(DequantSeverity.ERROR, strict.severity) + assertTrue(strict.profile.name.contains("strict")) + val failure = assertFailsWith { strict.requireAcceptable() } + assertTrue(failure.message!!.contains("over the 5.0%"), failure.message!!) + } + + // --- picking a profile --------------------------------------------------------------------- + + @Test + fun aDevicePicksItsOwnProfile() { + fun device(ramMb: Long) = DeviceMemory( + totalRamBytes = ramMb * mb, availableRamBytes = ramMb * mb / 2, + heapMaxBytes = 512 * mb, + ) + assertEquals(PlannerProfile.MOBILE_2GB, PlannerProfile.forDevice(device(2048))) + assertEquals(PlannerProfile.MOBILE_2GB, PlannerProfile.forDevice(device(3072)), "3 GB phones are still phones") + assertEquals(PlannerProfile.DESKTOP, PlannerProfile.forDevice(device(16384))) + } + + @Test + fun theRenderedPlanNamesTheProfileAndItsDecisions() { + val text = PlannerProfile.MOBILE_2GB.plan(input(ctx = 8192), availableBytes = 2048 * mb).render() + assertTrue(text.contains("profile mobile-2gb"), text) + assertTrue(text.contains("prefill 256"), text) + assertTrue(text.contains("off-heap ≥ 256"), text) + assertTrue(text.contains("KV auto-quantize over 80.0%"), text) + assertTrue(text.contains("note: KV cache switched"), text) + assertTrue(text.contains("note: weights are counted resident and mapped"), text) + } +}