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) + } +}