From dc6bb88f9ad092d5003a262b35ab208f1f319465 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 11:38:03 +0200 Subject: [PATCH] feat(android): mapped weights as the Android loader configuration, and a two-pool fit check before load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1038 (SKEEP-003 P5, S2.5, SKEEP-002, #921, #922). A phone has two memory pools and the planner totalled one number. The ART managed heap is hard-capped per app (256 MB, 512 MB with largeHeap) and holds every Kotlin array; mapped weights live in file-backed pages that do not count against that cap at all and compete for physical RAM instead, which the OS reclaims rather than killing the app. A single total cannot tell a model that will not fit from one that fits perfectly well as long as its weights are mapped. - `DeviceMemory` (RAM, the ART cap, the OS's own low-memory threshold) and `MemoryPlan.fitOn(device, weightsMapped)`: managed heap carries KV + forward + headroom, plus the weights only when they are not mapped; physical RAM carries everything, mapped pages included — they are evictable, not free. The RAM budget stays above the threshold the device itself declares. A failing fit names the pool that ran out and prices its advice against the actual shortfall, with "load with staging = MAPPED" first when that is what is missing. - `AndroidGguf`: `loader()` is the Android configuration of #1037's pipeline — `staging = MAPPED` by default; `deviceMemory(context)` reads `ActivityManager.MemoryInfo` plus `Runtime.maxMemory()`; `fits()` answers from the GGUF header before a byte of payload is read. - CI gains an `android` leg (`testAndroidHostTest`), and so does scripts/pr-gate.sh. Those host tests compile against androidMain, so nothing else in the gate proved that code even builds — `assemble` compiles it and runs nothing. SKEEP-002 stays Draft, with an implementation-status section recording what landed and what does not: packed weights still reach the heap, because the packed matmul SPI takes ByteArrays and a buffer-aware kernel needs the byte-order contract of #973 settled first. So the ceiling is lifted for dense checkpoints, not yet for a Q4_K_M one — which is exactly the M2-A5 criterion that stays open, together with the device measurements this repository has no hardware to make. DeviceFitTest (8 cases, every target): 640 MB of weights cannot live under a 512 MB cap and can when mapped; mapped pages still need RAM; the RAM budget respects the device's kill threshold; the first advice is to map; a fitting plan has no blocking pool. AndroidGgufLoadingHostTest (3 cases, Android compilation): the default loader maps, staging does not change the numbers, and the header-derived fit differs between mapped and heap by exactly the weight bytes. Gate: scripts/pr-gate.sh — all legs passed, including the new Android leg. Co-Authored-By: Claude Fable 5 --- .github/workflows/build.yml | 6 + .../002-android-offheap-tensor-storage.adoc | 38 +++++ scripts/pr-gate.sh | 6 + skainet-io/skainet-io-gguf/build.gradle.kts | 9 ++ .../io/gguf/AndroidGgufLoadingHostTest.kt | 145 ++++++++++++++++++ .../sk/ainet/io/gguf/AndroidGgufLoading.kt | 97 ++++++++++++ .../api/jvm/skainet-lang-core.api | 74 +++++++++ .../sk/ainet/lang/memory/plan/DeviceFit.kt | 139 +++++++++++++++++ .../ainet/lang/memory/plan/DeviceFitTest.kt | 125 +++++++++++++++ 9 files changed, 639 insertions(+) create mode 100644 skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/AndroidGgufLoadingHostTest.kt create mode 100644 skainet-io/skainet-io-gguf/src/androidMain/kotlin/sk/ainet/io/gguf/AndroidGgufLoading.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/DeviceFit.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/DeviceFitTest.kt diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0e2ee4c36..19c836007 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -41,6 +41,12 @@ jobs: tasks: verifyNpmPins jsTest wasmJsTest wasmWasiTest - name: native tasks: linuxX64Test + # Host-side (JVM) unit tests of the *Android* compilations: the mmap weight path (#921) + # and the Android loading facade + device fit check (#1038, SKEEP-002). They compile + # against androidMain, so no other leg proves that code builds — assemble compiles it but + # runs nothing. + - name: android + tasks: testAndroidHostTest # golden-parity: the SKEEP-003 packed-encoding gate (#1005). Bit-identical decode / # scalar-kernel / TurboQuant digests on JVM and Kotlin/Native plus the dispatch parity # tests, and the binary-compatibility check (apiCheck) that the contributing docs require diff --git a/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc b/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc index 964170c16..8f0364c4c 100644 --- a/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc +++ b/docs/modules/skeep/pages/002-android-offheap-tensor-storage.adoc @@ -232,6 +232,44 @@ keeps working unchanged and the overload is purely opt-in. Expected consumers: (packed quant blocks) and a synthesized SafeTensors file (dense F32/F16/BF16) produce bit-identical tensors under either placement. +== Implementation status (2026-08-24) + +Tracked by SKEEP-003's M2 slice +https://github.com/SKaiNET-developers/SKaiNET/issues/1038[#1038]. This SKEEP stays *Draft*: the +mechanism is in `develop`, the acceptance criteria that need a physical device are not met yet, and +one of them cannot be met until an unrelated contract is fixed. + +Landed: + +* *Phase 1 and 2* — `JvmMappedMemoryChunk` / `MappedRandomAccessSource` and a readable + `BufferHandle.FileBacked` with `JvmFileBackedResolver`, shared between the JVM and Android + compilations, covered by `androidHostTest` + (https://github.com/SKaiNET-developers/SKaiNET/issues/921[#921], + https://github.com/SKaiNET-developers/SKaiNET/issues/922[#922]). +* *Phase 3* — mapped loading is a *configuration of the ordinary loader* rather than a separate + helper: `StreamingGgufParametersLoader(staging = StagingPolicy.MAPPED)` + (https://github.com/SKaiNET-developers/SKaiNET/issues/1037[#1037]), with `AndroidGguf.loader()` + making it the Android default. +* *The fit check this SKEEP did not have* — `MemoryPlan.fitOn(DeviceMemory, weightsMapped)` treats + a phone as two pools (the ART cap and physical RAM) instead of one total, and + `AndroidGguf.fits(context, path, ctx)` answers from the GGUF header before a byte of payload is + read, naming the pool that runs out and what to do about it. + +Not landed, and why: + +* *Packed weights still reach the managed heap.* Mapped staging serves dense F32 tensors as + zero-heap views, but every quantized tensor is still handed to its kernel as a `ByteArray`, + because the packed matmul SPI takes arrays. That is Phase 4, and it is blocked on + https://github.com/SKaiNET-developers/SKaiNET/issues/973[#973] — the packed-quant byte-order + contract — since a buffer-aware kernel first needs an unambiguous answer to *which* byte order it + is reading. Until then the heap ceiling is lifted for dense checkpoints, not for a Q4_K_M one, + which is exactly the acceptance criterion below that remains open. +* *The device numbers.* "≤ 40 MB managed heap for SmolLM2-135M Q8_0" and "a ~600 MB Q4_K model + loads on a 256 MB heap" are measurements on a physical device; the repository's CI has none. They + belong to the `skainet-decode` sample in SKaiNET-transformers, which owns a model and a device + lane. `androidHostTest` covers what can be proven without hardware: the Android compilation + builds, maps, loads, and produces bit-identical tensors under either staging. + == Risks * *Page-fault latency.* First-touch of cold pages during decode adds jitter. diff --git a/scripts/pr-gate.sh b/scripts/pr-gate.sh index 67abe64c5..49e0d123a 100755 --- a/scripts/pr-gate.sh +++ b/scripts/pr-gate.sh @@ -56,6 +56,12 @@ step "assemble" step "Java consumer API tests" "${GRADLE[@]}" :skainet-test:skainet-test-java: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. +step "Android host tests" +"${GRADLE[@]}" testAndroidHostTest + if [[ "$mode" == "--bench" ]]; then step "benchmarks (compare against the committed baseline before/after)" "${GRADLE[@]}" :skainet-lang:skainet-lang-core:jvmBenchmark diff --git a/skainet-io/skainet-io-gguf/build.gradle.kts b/skainet-io/skainet-io-gguf/build.gradle.kts index 0936f979c..8fd5c2bf4 100644 --- a/skainet-io/skainet-io-gguf/build.gradle.kts +++ b/skainet-io/skainet-io-gguf/build.gradle.kts @@ -86,5 +86,14 @@ kotlin { implementation(libs.kotlinx.coroutines.test) } } + + // Host-side tests of the *Android* compilation (#1038): they drive the suspending loader, + // so they need coroutines like jvmTest does. The GGUF fixtures they write are their own — + // jvmTest's SyntheticGguf is not visible from this compilation. + getByName("androidHostTest") { + dependencies { + implementation(libs.kotlinx.coroutines) + } + } } } 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 new file mode 100644 index 000000000..bf40b3d1b --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/androidHostTest/kotlin/sk/ainet/io/gguf/AndroidGgufLoadingHostTest.kt @@ -0,0 +1,145 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +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.tensor.Tensor +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.MmapFloatTensorData +import sk.ainet.lang.types.FP32 +import java.io.File +import java.io.RandomAccessFile +import java.nio.ByteBuffer +import java.nio.ByteOrder +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * Host-side test of the *Android compilation* (#1038, SKEEP-002): this source set compiles against + * androidMain, so it proves the Android loading facade builds and behaves on the Android variant — + * mapped staging by default, and a fit check that answers before the load rather than during it. + * + * `AndroidGguf.deviceMemory(context)` needs a real `Context` and belongs to the instrumented smoke + * test; everything downstream of it takes a [DeviceMemory] so it can be checked here. + */ +@OptIn(ExperimentalMemoryApi::class) +class AndroidGgufLoadingHostTest { + + private val mb = 1024L * 1024L + + /** + * A one-tensor GGUF v3 file with a known F32 payload. Written here rather than reused from + * jvmTest's `SyntheticGguf`, which this compilation cannot see. + */ + private fun model(elements: Int = 4096): File { + val file = File.createTempFile("android-gguf-", ".gguf") + file.deleteOnExit() + val head = ByteBuffer.allocate(4096).order(ByteOrder.LITTLE_ENDIAN) + head.putInt(0x46554747) // "GGUF" + head.putInt(3) // version + head.putLong(1) // tensor count + head.putLong(1) // kv count + val key = "general.architecture".encodeToByteArray() + head.putLong(key.size.toLong()); head.put(key) + head.putInt(GGUFValueType.STRING.value) + val value = "test".encodeToByteArray() + head.putLong(value.size.toLong()); head.put(value) + val name = "w_f32".encodeToByteArray() + head.putLong(name.size.toLong()); head.put(name) + head.putInt(1) // rank + head.putLong(elements.toLong()) + head.putInt(GGMLQuantizationType.F32.value) + head.putLong(0L) // data offset + val padding = (32 - (head.position() % 32)) % 32 + repeat(padding) { head.put(0) } + + RandomAccessFile(file, "rw").use { raf -> + raf.write(head.array(), 0, head.position()) + val payload = ByteBuffer.allocate(elements * 4).order(ByteOrder.LITTLE_ENDIAN) + repeat(elements) { payload.putFloat(it * 0.5f) } + raf.write(payload.array()) + } + return file + } + + private fun load(loader: StreamingGgufParametersLoader): Map> { + val ctx = DefaultDataExecutionContext() + val out = LinkedHashMap>() + runBlocking { loader.load(ctx, FP32::class) { name, t -> out[name] = t } } + return out + } + + @Test + fun `the android loader maps weights by default`() { + val f = model() + try { + val mapped = load(AndroidGguf.loader(f.absolutePath)) + assertTrue( + mapped.getValue("w_f32").data is MmapFloatTensorData<*>, + "dense F32 must come from file-backed pages on Android, got ${mapped.getValue("w_f32").data::class.simpleName}", + ) + // and the heap path is still reachable, producing the same numbers + val onHeap = load(AndroidGguf.loader(f.absolutePath, staging = StagingPolicy.HEAP)) + assertTrue(onHeap.getValue("w_f32").data is FloatArrayTensorData<*>) + assertContentEquals( + onHeap.getValue("w_f32").data.copyToFloatArray(), + mapped.getValue("w_f32").data.copyToFloatArray(), + "staging must not change the numbers", + ) + val values = mapped.getValue("w_f32").data.copyToFloatArray() + assertEquals(0f, values[0]); assertEquals(0.5f, values[1]); assertEquals(2047.5f, values[4095]) + } finally { + f.delete() + } + } + + @Test + fun `the fit check reads the plan from the header and answers before loading`() { + val f = model() + try { + val plentiful = DeviceMemory( + totalRamBytes = 4096 * mb, availableRamBytes = 2048 * mb, + heapMaxBytes = 512 * mb, heapUsedBytes = 32 * mb, lowMemoryThresholdBytes = 180 * mb, + ) + val fit = AndroidGguf.fits(plentiful, f.absolutePath, ctx = 128) + assertTrue(fit.fits, fit.render()) + assertTrue(fit.weightsMapped, "the Android default is mapped weights") + assertTrue(fit.plan.weightsBytes > 0, "the plan comes from the header") + + // a phone with almost nothing left says so, and says which pool ran out + val squeezed = plentiful.copy(availableRamBytes = 190 * mb, heapMaxBytes = 8 * mb, heapUsedBytes = 7 * mb) + val tight = AndroidGguf.fits(squeezed, f.absolutePath, ctx = 4096) + assertFalse(tight.fits, tight.render()) + assertEquals("managed heap", tight.blockingPool) + assertTrue(tight.suggestions.isNotEmpty(), "a failing fit must say what to do") + } finally { + f.delete() + } + } + + @Test + fun `an unmapped load is charged for its weights, a mapped one is not`() { + val f = model() + try { + val device = DeviceMemory( + totalRamBytes = 2048 * mb, availableRamBytes = 900 * mb, + heapMaxBytes = 512 * mb, heapUsedBytes = 40 * mb, lowMemoryThresholdBytes = 180 * mb, + ) + val mapped = AndroidGguf.fits(device, f.absolutePath, ctx = 512, weightsMapped = true) + val heap = AndroidGguf.fits(device, f.absolutePath, ctx = 512, weightsMapped = false) + assertEquals( + mapped.plan.weightsBytes, + heap.heap.neededBytes - mapped.heap.neededBytes, + "the difference between the two is exactly the weights", + ) + } finally { + f.delete() + } + } +} 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 new file mode 100644 index 000000000..e6d546448 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/androidMain/kotlin/sk/ainet/io/gguf/AndroidGgufLoading.kt @@ -0,0 +1,97 @@ +package sk.ainet.io.gguf + +import android.app.ActivityManager +import android.content.Context +import sk.ainet.io.RandomAccessSource +import sk.ainet.io.openRandomAccessSource +import sk.ainet.io.model.QuantPolicy +import sk.ainet.io.model.StagingPolicy +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.Budget +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.fitOn + +/** + * Loading a GGUF on Android: mapped weights by default, and a fit check *before* the load + * (SKEEP-002, #921, #922, #1038). + * + * The managed heap is the binding constraint on a phone — hard-capped at 256 MB (512 MB with + * `largeHeap`) no matter how much RAM the device has — so the Android configuration of the loader + * is `staging = MAPPED`: weights come from file-backed pages the OS pages in on demand and evicts + * under pressure, and never count against the cap. + * + * What is *not* solved yet: packed (quantized) tensors still arrive as heap arrays, because the + * packed kernels take `ByteArray`s until the view contract of #973 lands. Mapping therefore lifts + * the ceiling for dense-F32 weights today, and for a Q4_K_M checkpoint only once #973 does. The + * fit check tells you which of the two pools you are about to run out of, rather than letting the + * app find out by being killed. + */ +@OptIn(ExperimentalMemoryApi::class) +public object AndroidGguf { + + /** + * The loader Android should use: positional reads for the metadata, mapped pages for tensor + * payloads. [quantPolicy] is the caller's choice as usual; [staging] defaults to + * [StagingPolicy.MAPPED] and is a parameter only so a test or a benchmark can ask for the + * heap path explicitly. + */ + public fun loader( + filePath: String, + quantPolicy: QuantPolicy = QuantPolicy.NATIVE_OPTIMIZED, + staging: StagingPolicy = StagingPolicy.MAPPED, + onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + ): StreamingGgufParametersLoader = StreamingGgufParametersLoader( + sourceProvider = { openSource(filePath) }, + onProgress = onProgress, + quantPolicy = quantPolicy, + staging = staging, + ) + + /** + * What this device has to offer: `ActivityManager.MemoryInfo` for physical RAM plus the ART + * heap cap, which is what actually stops a model from loading. + */ + public fun deviceMemory(context: Context): DeviceMemory { + val am = context.getSystemService(Context.ACTIVITY_SERVICE) as ActivityManager + val info = ActivityManager.MemoryInfo() + am.getMemoryInfo(info) + val runtime = Runtime.getRuntime() + return DeviceMemory( + totalRamBytes = info.totalMem, + availableRamBytes = info.availMem, + heapMaxBytes = runtime.maxMemory(), + heapUsedBytes = runtime.totalMemory() - runtime.freeMemory(), + lowMemory = info.lowMemory, + lowMemoryThresholdBytes = info.threshold, + ) + } + + /** + * The plan a GGUF's *header* predicts at [ctx] — tensor table and metadata only, no payload + * (M0-F1), so this costs a few kilobytes and a couple of reads. + */ + public fun plan(filePath: String, ctx: Int, budget: Budget? = null): MemoryPlan = + openSource(filePath).use { source -> + MemoryPlans.plan(StreamingGGUFReader.open(source).planInput(ctx), budget) + } + + private fun openSource(filePath: String): RandomAccessSource = + openRandomAccessSource(filePath) + ?: throw IllegalArgumentException("Cannot open for random access: $filePath") + + /** + * Will this model load on this device? Checks the header-derived plan against both pools — + * managed heap and physical RAM — before a byte of payload is read. + * + * @param weightsMapped whether the load will use [StagingPolicy.MAPPED] (what [loader] does) + */ + public fun fits(context: Context, filePath: String, ctx: Int, weightsMapped: Boolean = true): DeviceFit = + fits(deviceMemory(context), filePath, ctx, weightsMapped) + + /** [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) +} 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 a7e879ab3..6ea9024f2 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,63 @@ 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/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; + public final fun component2 ()Lsk/ainet/lang/memory/plan/DeviceMemory; + public final fun component3 ()Z + public final fun component4 ()Lsk/ainet/lang/memory/plan/PoolFit; + public final fun component5 ()Lsk/ainet/lang/memory/plan/PoolFit; + public final fun component6 ()Ljava/util/List; + public final fun copy (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;)Lsk/ainet/lang/memory/plan/DeviceFit; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/DeviceFit;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;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/DeviceFit; + public fun equals (Ljava/lang/Object;)Z + public final fun getBlockingPool ()Ljava/lang/String; + public final fun getDevice ()Lsk/ainet/lang/memory/plan/DeviceMemory; + public final fun getFits ()Z + public final fun getHeap ()Lsk/ainet/lang/memory/plan/PoolFit; + public final fun getPlan ()Lsk/ainet/lang/memory/plan/MemoryPlan; + public final fun getRam ()Lsk/ainet/lang/memory/plan/PoolFit; + public final fun getSuggestions ()Ljava/util/List; + public final fun getWeightsMapped ()Z + public fun hashCode ()I + public final fun render ()Ljava/lang/String; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/DeviceFitKt { + public static final fun fitOn (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/DeviceMemory;Z)Lsk/ainet/lang/memory/plan/DeviceFit; +} + +public final class sk/ainet/lang/memory/plan/DeviceMemory { + public static final field Companion Lsk/ainet/lang/memory/plan/DeviceMemory$Companion; + public static final field RAM_RESERVE_FLOOR J + public fun (JJJJZJ)V + public synthetic fun (JJJJZJILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()J + public final fun component2 ()J + public final fun component3 ()J + public final fun component4 ()J + public final fun component5 ()Z + public final fun component6 ()J + public final fun copy (JJJJZJ)Lsk/ainet/lang/memory/plan/DeviceMemory; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/DeviceMemory;JJJJZJILjava/lang/Object;)Lsk/ainet/lang/memory/plan/DeviceMemory; + public fun equals (Ljava/lang/Object;)Z + public final fun getAvailableRamBytes ()J + public final fun getHeapFreeBytes ()J + public final fun getHeapMaxBytes ()J + public final fun getHeapUsedBytes ()J + public final fun getLowMemory ()Z + public final fun getLowMemoryThresholdBytes ()J + public final fun getTotalRamBytes ()J + public final fun getUsableRamBytes ()J + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/DeviceMemory$Companion { +} + public final class sk/ainet/lang/memory/plan/KvCacheMode : java/lang/Enum { public static final field BF16 Lsk/ainet/lang/memory/plan/KvCacheMode; public static final field FP32 Lsk/ainet/lang/memory/plan/KvCacheMode; @@ -1490,6 +1547,23 @@ public final class sk/ainet/lang/memory/plan/PlanVsActualLine { public final fun withinTolerance (D)Z } +public final class sk/ainet/lang/memory/plan/PoolFit { + public fun (Ljava/lang/String;JJ)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()J + public final fun component3 ()J + public final fun copy (Ljava/lang/String;JJ)Lsk/ainet/lang/memory/plan/PoolFit; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PoolFit;Ljava/lang/String;JJILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PoolFit; + public fun equals (Ljava/lang/Object;)Z + public final fun getBudgetBytes ()J + public final fun getFits ()Z + public final fun getHeadroomBytes ()J + public final fun getName ()Ljava/lang/String; + public final fun getNeededBytes ()J + public fun hashCode ()I + 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/DeviceFit.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/DeviceFit.kt new file mode 100644 index 000000000..2f595c65b --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/DeviceFit.kt @@ -0,0 +1,139 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** + * What a device has to offer, as **two pools** rather than one (SKEEP-002, #1038). + * + * On Android these are genuinely different resources: the ART managed heap is hard-capped per app + * (256 MB by default, 512 MB with `largeHeap`) and is where every Kotlin array lives, while mapped + * weights live in file-backed pages that do not count against that cap at all — they compete for + * physical RAM, which the OS reclaims under pressure instead of killing the app. A planner that + * totals one number cannot tell a model that will not fit from one that fits perfectly well as + * long as its weights are mapped. + * + * @property totalRamBytes physical RAM (`ActivityManager.MemoryInfo.totalMem`) + * @property availableRamBytes RAM the OS says is available now (`MemoryInfo.availMem`) + * @property lowMemory the OS is already under pressure (`MemoryInfo.lowMemory`) + * @property lowMemoryThresholdBytes below this the OS starts killing processes (`MemoryInfo.threshold`) + * @property heapMaxBytes the managed-heap cap (`Runtime.maxMemory()`) + * @property heapUsedBytes managed heap in use right now + */ +@ExperimentalMemoryApi +public data class DeviceMemory( + val totalRamBytes: Long, + val availableRamBytes: Long, + val heapMaxBytes: Long, + val heapUsedBytes: Long = 0L, + val lowMemory: Boolean = false, + val lowMemoryThresholdBytes: Long = 0L, +) { + init { + require(totalRamBytes >= 0 && availableRamBytes >= 0) { "RAM figures must be non-negative" } + require(heapMaxBytes > 0) { "heapMaxBytes must be > 0" } + } + + /** Managed heap still available to allocate into. */ + public val heapFreeBytes: Long get() = (heapMaxBytes - heapUsedBytes).coerceAtLeast(0L) + + /** + * RAM that may be used without pushing the OS under its own low-memory threshold — the reserve + * the device itself declares, floored at [RAM_RESERVE_FLOOR] for devices that report none. + */ + public val usableRamBytes: Long + get() = (availableRamBytes - maxOf(lowMemoryThresholdBytes, RAM_RESERVE_FLOOR)).coerceAtLeast(0L) + + public companion object { + /** Reserve when the device reports no low-memory threshold. */ + public const val RAM_RESERVE_FLOOR: Long = 128L * MiB + } +} + +/** One resource pool of a [DeviceFit]: what the plan needs from it, and what it has. */ +@ExperimentalMemoryApi +public data class PoolFit(val name: String, val neededBytes: Long, val budgetBytes: Long) { + public val fits: Boolean get() = neededBytes <= budgetBytes + /** Bytes left over (negative when it does not fit). */ + public val headroomBytes: Long get() = budgetBytes - neededBytes +} + +/** + * A [MemoryPlan] checked against a real device, pool by pool (M2-A5): the answer to "will this + * model load on this phone", and when it will not, which pool ran out and what to do about it. + * + * @property weightsMapped whether the weights are loaded through `StagingPolicy.MAPPED`, i.e. from + * file-backed pages that never count against the managed heap. + */ +@ExperimentalMemoryApi +public data class DeviceFit( + val plan: MemoryPlan, + val device: DeviceMemory, + val weightsMapped: Boolean, + val heap: PoolFit, + val ram: PoolFit, + val suggestions: List, +) { + /** True only when *both* pools have room. */ + public val fits: Boolean get() = heap.fits && ram.fits + + /** The pool that ran out, or `null` when the plan fits. */ + public val blockingPool: String? + get() = when { + !heap.fits -> heap.name + !ram.fits -> ram.name + else -> null + } + + public fun render(): String = buildString { + append(plan.input.modelName); append(" · ctx "); append(plan.input.ctx) + append(if (weightsMapped) " · weights mapped\n" else " · weights on the heap\n") + for (p in listOf(heap, ram)) { + append(" "); append(p.name.padEnd(14)) + append(MemoryPlans.formatBytes(p.neededBytes).padStart(10)) + append(" of "); append(MemoryPlans.formatBytes(p.budgetBytes).padStart(10)) + append(if (p.fits) " ✔" else " ✘ short by ${MemoryPlans.formatBytes(-p.headroomBytes)}") + append('\n') + } + if (device.lowMemory) append(" device reports low memory — the OS is already reclaiming\n") + if (!fits && suggestions.isNotEmpty()) { + append(" suggestions: ") + append(suggestions.joinToString(" · ") { "${it.text} (−${MemoryPlans.formatBytes(it.savesBytes)})" }) + append('\n') + } + } +} + +/** + * Check [plan] against [device]'s two pools (SKEEP-002, M2-A5). + * + * The managed heap carries the KV cache, the forward slab and the plan's heap headroom — plus the + * weights when they are *not* mapped, which is exactly what makes a 600 MB Q4 model impossible + * under a 512 MB cap and unremarkable when mapped. Physical RAM carries everything, mapped pages + * included: they are evictable, not free. + * + * @param weightsMapped weights come from file-backed pages (`StagingPolicy.MAPPED`) + */ +@ExperimentalMemoryApi +public fun MemoryPlan.fitOn(device: DeviceMemory, weightsMapped: Boolean): DeviceFit { + val heapNeeded = kvBytes + forwardBytes + headroomBytes + if (weightsMapped) 0L else weightsBytes + val heap = PoolFit("managed heap", heapNeeded, device.heapFreeBytes) + val ram = PoolFit("device RAM", totalBytes, device.usableRamBytes) + + val suggestions = ArrayList() + if (!heap.fits && !weightsMapped) { + suggestions += Suggestion( + "load with staging = MAPPED (weights move to file-backed pages, off the managed heap)", + weightsBytes, + ) + } + // The plan's own advice (smaller ctx, quantized KV, a smaller model) is priced against a + // budget, and the budget here is "whatever the tighter pool is short by": a plan that misses + // the heap cap by 200 MB has to save 200 MB, wherever those bytes were going to sit. + val deficit = maxOf(-heap.headroomBytes, -ram.headroomBytes) + if (deficit > 0) { + val target = (totalBytes - deficit).coerceAtLeast(0L) + val pool = if (!heap.fits) heap.name else ram.name + suggestions += copy(budget = Budget(target, "$pool on this device")).suggestions() + } + return DeviceFit(this, device, weightsMapped, heap, ram, suggestions) +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/DeviceFitTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/DeviceFitTest.kt new file mode 100644 index 000000000..027cc2c72 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/DeviceFitTest.kt @@ -0,0 +1,125 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** + * #1038 (SKEEP-002, M2-A5): a phone has **two** memory pools, and the plan has to be checked + * against both. + * + * The managed heap is hard-capped per app and holds every Kotlin array; physical RAM holds + * everything including mapped pages. A 600 MB Q4 model is impossible under a 512 MB cap and + * unremarkable when its weights are mapped — one total cannot express that, so this asserts the + * two-pool arithmetic and the advice it produces. + */ +@OptIn(ExperimentalMemoryApi::class) +class DeviceFitTest { + + private val mb = 1024L * 1024L + + /** A Llama-1B-shaped Q4_K model: ~640 MB of weights. */ + private fun plan(ctx: Int = 2048, weightsMb: Long = 640): MemoryPlan { + val f = Format(FP32, TensorEncoding.Q4_K) + val elementsPerTensor = weightsMb * mb / 144 * 256 // Q4_K: 144 bytes per 256 elements + val tensors = listOf(PlanTensor("model.weight", null, f, elementsPerTensor, weightsMb * mb)) + val geometry = ModelGeometry( + layers = 16, heads = 32, kvHeads = 8, headDim = 64, + embeddingLength = 2048, feedForwardLength = 5632, vocabSize = 32000, + ) + return MemoryPlans.plan(PlanInput("llama-1b", "llama", tensors, geometry, ctx)) + } + + /** A 2 GB phone: 512 MB largeHeap cap, ~900 MB RAM free, OS kills below 180 MB. */ + private fun phone(heapMaxMb: Long = 512, heapUsedMb: Long = 40, availMb: Long = 900) = DeviceMemory( + totalRamBytes = 2048 * mb, + availableRamBytes = availMb * mb, + heapMaxBytes = heapMaxMb * mb, + heapUsedBytes = heapUsedMb * mb, + lowMemoryThresholdBytes = 180 * mb, + ) + + @Test + fun heapWeightsDoNotFitUnderTheArtCapButMappedOnesDo() { + val p = plan() + val onHeap = p.fitOn(phone(), weightsMapped = false) + assertFalse(onHeap.fits, "640 MB of weights cannot live under a 512 MB cap:\n${onHeap.render()}") + assertEquals("managed heap", onHeap.blockingPool) + assertTrue(onHeap.heap.neededBytes >= p.weightsBytes, "unmapped weights are charged to the heap") + + val mapped = p.fitOn(phone(), weightsMapped = true) + assertTrue(mapped.heap.fits, "mapped weights never touch the managed heap:\n${mapped.render()}") + assertEquals(p.kvBytes + p.forwardBytes + p.headroomBytes, mapped.heap.neededBytes) + } + + @Test + fun mappedPagesStillNeedPhysicalRam() { + val p = plan() + // plenty of heap, but the device has almost no free RAM: mapping is not free memory + val squeezed = p.fitOn(phone(availMb = 300), weightsMapped = true) + assertFalse(squeezed.fits, squeezed.render()) + assertEquals("device RAM", squeezed.blockingPool) + assertEquals(p.totalBytes, squeezed.ram.neededBytes, "RAM carries everything, mapped included") + } + + @Test + fun theRamBudgetKeepsTheDeviceAboveItsOwnLowMemoryThreshold() { + val p = plan(weightsMb = 100) + val device = phone(availMb = 600) + val fit = p.fitOn(device, weightsMapped = true) + assertEquals((600 - 180) * mb, fit.ram.budgetBytes, "available minus the OS's own kill threshold") + + // a device that reports no threshold still keeps a floor + val noThreshold = device.copy(lowMemoryThresholdBytes = 0) + assertEquals(600 * mb - DeviceMemory.RAM_RESERVE_FLOOR, noThreshold.usableRamBytes) + } + + @Test + fun theFirstAdviceIsToMapTheWeights() { + val fit = plan().fitOn(phone(), weightsMapped = false) + val first = fit.suggestions.first() + assertTrue(first.text.contains("MAPPED"), "expected mapping advice first, got '${first.text}'") + assertEquals(fit.plan.weightsBytes, first.savesBytes, "it saves exactly the weights") + assertTrue(fit.suggestions.size > 1, "and the plan's own suggestions follow: ${fit.suggestions.map { it.text }}") + } + + @Test + fun aPlanThatFitsHasNoBlockingPoolAndNoAdvice() { + val fit = plan(ctx = 512, weightsMb = 120).fitOn(phone(), weightsMapped = true) + assertTrue(fit.fits, fit.render()) + assertNull(fit.blockingPool) + assertTrue(fit.suggestions.isEmpty()) + assertTrue(fit.heap.headroomBytes > 0 && fit.ram.headroomBytes > 0) + } + + @Test + fun theRenderedTableNamesThePoolAndTheShortfall() { + val text = plan().fitOn(phone(), weightsMapped = false).render() + assertTrue(text.contains("weights on the heap"), text) + assertTrue(text.contains("managed heap"), text) + assertTrue(text.contains("device RAM"), text) + assertTrue(text.contains("short by"), text) + assertTrue(plan(ctx = 512, weightsMb = 120).fitOn(phone(), weightsMapped = true).render().contains("weights mapped")) + } + + @Test + fun aDeviceUnderPressureSaysSo() { + val fit = plan(weightsMb = 100).fitOn(phone().copy(lowMemory = true), weightsMapped = true) + assertTrue(fit.render().contains("low memory"), fit.render()) + } + + @Test + fun heapUseCountsAgainstTheCap() { + val p = plan(weightsMb = 100) + val fresh = p.fitOn(phone(heapUsedMb = 0), weightsMapped = false) + val busy = p.fitOn(phone(heapUsedMb = 400), weightsMapped = false) + assertEquals(fresh.heap.neededBytes, busy.heap.neededBytes) + assertTrue(fresh.heap.budgetBytes > busy.heap.budgetBytes, "an app already holding 400 MB has less room") + } +}