From c1f6718b956c990f2b40e1b2d027aef0c2ea3e69 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 10:36:29 +0200 Subject: [PATCH] =?UTF-8?q?feat(plan):=20AllocationResolver=20=E2=80=94=20?= =?UTF-8?q?one=20resolver=20decides=20domain=20and=20scope?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Placement joins WeightForm as a resolved decision (#1133's answer): AllocationResolver is a pure function of what will be held (the resolved form), what the profile says (domainFor and its threshold) and what the platform can do (StorageCapabilities, injectable so a test can resolve for a platform it is not running on). PlanTensor.allocation stops hardcoding MMAP_FILE/MODEL — a spec that claimed every weight was mapped even on a platform that cannot map, and even for a dequantized copy that no longer matches the file. Mapping now requires all three: the form asks MAPPED, the platform can map, and the bytes really are the file's bytes; everything else falls to the profile's heap/off-heap threshold over the bytes actually held. AllocationResolver.explain() renders the decision with its reason, so a plan can say where every tensor landed and why. Closes #1143. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 35 +++- .../lang/memory/plan/AllocationResolver.kt | 131 +++++++++++++++ .../sk/ainet/lang/memory/plan/MemoryPlan.kt | 14 +- .../memory/plan/AllocationResolverTest.kt | 149 ++++++++++++++++++ .../ainet/lang/memory/plan/MemoryPlanTest.kt | 11 +- .../lang/memory/plan/PlannerProfileTest.kt | 8 + 6 files changed, 339 insertions(+), 9 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/AllocationResolverTest.kt 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 2257c6923..a3a27bcab 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 @@ -1360,6 +1360,17 @@ public final class sk/ainet/lang/memory/plan/ActualMemory$Companion { public final fun from (Lsk/ainet/lang/memory/trace/RecordingTraceSink;)Lsk/ainet/lang/memory/plan/ActualMemory; } +public final class sk/ainet/lang/memory/plan/AllocationResolver { + public static final field INSTANCE Lsk/ainet/lang/memory/plan/AllocationResolver; + public final fun explain (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Ljava/lang/String; + public static synthetic fun explain$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Ljava/lang/String; + public final fun residentFormat (Lsk/ainet/lang/memory/plan/PlanTensor;)Lsk/ainet/lang/memory/Format; + public final fun resolve (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Lsk/ainet/lang/memory/AllocationSpec; + public static synthetic fun resolve$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; + public final fun resolveTransient (Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;Lsk/ainet/lang/memory/ScopeKind;)Lsk/ainet/lang/memory/AllocationSpec; + public static synthetic fun resolveTransient$default (Lsk/ainet/lang/memory/plan/AllocationResolver;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;Lsk/ainet/lang/memory/ScopeKind;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; +} + public final class sk/ainet/lang/memory/plan/Budget { public static final field Companion Lsk/ainet/lang/memory/plan/Budget$Companion; public static final field RESERVE_ANDROID_JVM J @@ -1649,6 +1660,8 @@ public final class sk/ainet/lang/memory/plan/PlanLine { public final class sk/ainet/lang/memory/plan/PlanTensor { public fun (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;)V public synthetic fun (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun allocation (Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;)Lsk/ainet/lang/memory/AllocationSpec; + public static synthetic fun allocation$default (Lsk/ainet/lang/memory/plan/PlanTensor;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/StorageCapabilities;ILjava/lang/Object;)Lsk/ainet/lang/memory/AllocationSpec; public final fun component1 ()Ljava/lang/String; public final fun component2 ()Lsk/ainet/lang/tensor/TensorId; public final fun component3 ()Lsk/ainet/lang/memory/Format; @@ -1658,7 +1671,6 @@ public final class sk/ainet/lang/memory/plan/PlanTensor { public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;)Lsk/ainet/lang/memory/plan/PlanTensor; public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanTensor;Ljava/lang/String;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/Format;JJLsk/ainet/lang/memory/plan/WeightForm;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanTensor; public fun equals (Ljava/lang/Object;)Z - public final fun getAllocation ()Lsk/ainet/lang/memory/AllocationSpec; public final fun getBytes ()J public final fun getElementCount ()J public final fun getForm ()Lsk/ainet/lang/memory/plan/WeightForm; @@ -1795,6 +1807,27 @@ public final class sk/ainet/lang/memory/plan/ProfiledPlan { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/plan/StorageCapabilities { + public static final field Companion Lsk/ainet/lang/memory/plan/StorageCapabilities$Companion; + public fun (ZZ)V + public synthetic fun (ZZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Z + public final fun component2 ()Z + public final fun copy (ZZ)Lsk/ainet/lang/memory/plan/StorageCapabilities; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/StorageCapabilities;ZZILjava/lang/Object;)Lsk/ainet/lang/memory/plan/StorageCapabilities; + public fun equals (Ljava/lang/Object;)Z + public final fun getSupportsMappedFiles ()Z + public final fun getSupportsOffHeap ()Z + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/StorageCapabilities$Companion { + public final fun current ()Lsk/ainet/lang/memory/plan/StorageCapabilities; + public final fun getFULL ()Lsk/ainet/lang/memory/plan/StorageCapabilities; + public final fun getHEAP_ONLY ()Lsk/ainet/lang/memory/plan/StorageCapabilities; +} + 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/AllocationResolver.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt new file mode 100644 index 000000000..3f28fbdf6 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/AllocationResolver.kt @@ -0,0 +1,131 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.AllocationSpec +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.PlatformStorage +import sk.ainet.lang.memory.ScopeKind +import sk.ainet.lang.tensor.storage.MemoryDomain + +/** + * What the running platform's storage can actually do — the third input of [AllocationResolver], + * separated from [PlatformStorage] so a test can resolve for a platform it is not running on. + */ +@ExperimentalMemoryApi +public data class StorageCapabilities( + val supportsMappedFiles: Boolean, + val supportsOffHeap: Boolean = true, +) { + public companion object { + /** The platform this code is running on. */ + public fun current(): StorageCapabilities = StorageCapabilities( + supportsMappedFiles = PlatformStorage.supportsMappedFiles, + supportsOffHeap = PlatformStorage.supports(MemoryDomain.HOST_OFFHEAP), + ) + + /** A JVM/native-class platform: everything works. */ + public val FULL: StorageCapabilities = StorageCapabilities(supportsMappedFiles = true) + + /** A browser-class platform: heap only, no mmap, no off-heap. */ + public val HEAP_ONLY: StorageCapabilities = StorageCapabilities(supportsMappedFiles = false, supportsOffHeap = false) + } +} + +/** + * Decides where a tensor's bytes belong — the placement counterpart of [WeightFormResolver] (#1143, + * closing the question #1133 asked). + * + * Same contract as the form resolver: a pure function of *what will be held* (the resolved + * [WeightForm]), *what the profile says* ([PlannerProfile.domainFor], its thresholds) and *what the + * platform can do* ([StorageCapabilities]). Consumers — the plan, the loader, a context — carry the + * result; none of them decide. Nothing here allocates. + */ +@ExperimentalMemoryApi +public object AllocationResolver { + + /** + * The allocation a weight needs once [PlanTensor.form] is honoured. + * + * Served from file-backed pages only when every condition holds: the form asks for + * [WeightResidency.MAPPED], the platform can map, and the bytes really are the file's bytes — + * a re-encoded ([EncodingRequest.DequantizeTo]/[EncodingRequest.RequantizeTo]) or re-ordered + * ([WeightByteOrder.KERNEL_FEED]) weight is a load-time copy, and a copy cannot be paged from + * the file it no longer matches. Everything else falls to [PlannerProfile.domainFor] over the + * bytes actually held, so a dequantized giant goes off-heap and a small bias stays on it. + */ + public fun resolve( + weight: PlanTensor, + profile: PlannerProfile, + platform: StorageCapabilities = StorageCapabilities.current(), + ): AllocationSpec { + val form = weight.form + val fileBytesAreTheBytes = form == null || + (form.encoding == EncodingRequest.KeepAsStored && form.order == WeightByteOrder.AS_STORED) + val wantsMapped = form?.residency == WeightResidency.MAPPED + val mapped = wantsMapped && platform.supportsMappedFiles && fileBytesAreTheBytes + val domain = if (mapped) MemoryDomain.MMAP_FILE else fallbackDomain(weight.residentBytes, profile, platform) + return AllocationSpec( + format = residentFormat(weight), + elementCount = weight.elementCount, + domain = domain, + scope = ScopeKind.MODEL, + mutable = false, + ) + } + + /** + * The allocation a transient (activation/scratch) tensor needs: never mapped, never + * model-lifetime — only the profile's heap/off-heap threshold and the caller's scope. + */ + public fun resolveTransient( + format: Format, + elementCount: Long, + profile: PlannerProfile, + platform: StorageCapabilities = StorageCapabilities.current(), + scope: ScopeKind = ScopeKind.FORWARD, + ): AllocationSpec { + val bytes = format.physicalBytes(elementCount) ?: (format.dtype.sizeInBytes.toLong() * elementCount) + return AllocationSpec(format, elementCount, fallbackDomain(bytes, profile, platform), scope, mutable = true) + } + + /** The [Format] the weight holds once its form is honoured — dense after a dequantization. */ + public fun residentFormat(weight: PlanTensor): Format = when (val request = weight.form?.encoding) { + null, EncodingRequest.KeepAsStored -> weight.format + is EncodingRequest.DequantizeTo -> Format.dense(request.dtype) + is EncodingRequest.RequantizeTo -> Format(weight.format.dtype, request.encoding) + } + + /** + * One line saying where [weight] lands and *why* — the transparency counterpart of + * [resolve], for plan renders and load-time traces. The decision is recomputed, so the + * explanation can never drift from what the resolver actually did. + */ + public fun explain( + weight: PlanTensor, + profile: PlannerProfile, + platform: StorageCapabilities = StorageCapabilities.current(), + ): String { + val spec = resolve(weight, profile, platform) + val form = weight.form + val why = when { + spec.domain == MemoryDomain.MMAP_FILE -> + "form asks MAPPED, platform can map, bytes are the file's bytes" + form?.residency == WeightResidency.MAPPED && !platform.supportsMappedFiles -> + "form asks MAPPED but this platform cannot map files" + form?.residency == WeightResidency.MAPPED && form.encoding != EncodingRequest.KeepAsStored -> + "form asks MAPPED but the weight is re-encoded at load — a copy cannot be paged from the file" + form?.residency == WeightResidency.MAPPED && form.order == WeightByteOrder.KERNEL_FEED -> + "form asks MAPPED but kernel-feed order is a load-time copy" + else -> + "resident ${MemoryPlans.formatBytes(weight.residentBytes)} vs off-heap threshold " + + MemoryPlans.formatBytes(profile.offHeapThresholdBytes) + + if (!platform.supportsOffHeap) " (no off-heap on this platform)" else "" + } + return "${weight.name}: ${spec.domain}/${spec.scope} — $why" + } + + private fun fallbackDomain(bytes: Long, profile: PlannerProfile, platform: StorageCapabilities): MemoryDomain { + val preferred = profile.domainFor(bytes) + return if (preferred == MemoryDomain.HOST_OFFHEAP && !platform.supportsOffHeap) MemoryDomain.HOST_HEAP else preferred + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt index 9c84165f9..d99b80d37 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/MemoryPlan.kt @@ -3,9 +3,7 @@ package sk.ainet.lang.memory.plan import sk.ainet.lang.memory.AllocationSpec import sk.ainet.lang.memory.ExperimentalMemoryApi import sk.ainet.lang.memory.Format -import sk.ainet.lang.memory.ScopeKind import sk.ainet.lang.tensor.TensorId -import sk.ainet.lang.tensor.storage.MemoryDomain import sk.ainet.lang.tensor.storage.TensorEncoding /** @@ -44,9 +42,15 @@ public data class PlanTensor( Format(format.dtype, request.encoding).physicalBytes(elementCount) ?: bytes } - /** The allocation this weight needs: mapped, model-lifetime, read-only. */ - val allocation: AllocationSpec - get() = AllocationSpec(format, elementCount, MemoryDomain.MMAP_FILE, ScopeKind.MODEL, mutable = false) + /** + * The allocation this weight needs — resolved, not assumed (#1143). The old form of this + * property hardcoded mapped/model-lifetime for every weight regardless of what [form] asked, + * what the profile said, or whether the platform could map at all. + */ + public fun allocation( + profile: PlannerProfile, + platform: StorageCapabilities = StorageCapabilities.current(), + ): AllocationSpec = AllocationResolver.resolve(this, profile, platform) } /** The transformer geometry the KV-cache and forward-slab estimates need (from the GGUF header). */ diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/AllocationResolverTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/AllocationResolverTest.kt new file mode 100644 index 000000000..9ca54e70a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/AllocationResolverTest.kt @@ -0,0 +1,149 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ScopeKind +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.assertFalse +import kotlin.test.assertNotEquals +import kotlin.test.assertTrue + +/** + * #1143 / #1133: placement is a resolver decision — (what will be held × profile × platform) in, + * an [sk.ainet.lang.memory.AllocationSpec] out. These tests pin the rules the old + * `PlanTensor.allocation` hardcode ignored. + */ +@OptIn(ExperimentalMemoryApi::class) +class AllocationResolverTest { + + private val q4k = Format(FP32, TensorEncoding.Q4_K) + + private fun weight( + elements: Long = 1L shl 20, // 1Mi elements: Q4_K packed ≈ 576 KiB, dense FP32 = 4 MiB + form: WeightForm?, + ) = PlanTensor( + name = "blk.0.attn_q.weight", + id = null, + format = q4k, + elementCount = elements, + bytes = q4k.physicalBytes(elements)!!, + form = form, + ) + + @Test + fun mappedKeptAsStoredWeightIsServedFromTheFile() { + val spec = AllocationResolver.resolve( + weight(form = WeightForm(residency = WeightResidency.MAPPED)), + PlannerProfile.MOBILE_2GB, + StorageCapabilities.FULL, + ) + assertEquals(MemoryDomain.MMAP_FILE, spec.domain) + assertEquals(ScopeKind.MODEL, spec.scope) + assertFalse(spec.mutable) + assertEquals(q4k, spec.format) + } + + @Test + fun unmappablePlatformFallsBackByTheProfileThreshold() { + val spec = AllocationResolver.resolve( + weight(form = WeightForm(residency = WeightResidency.MAPPED)), + PlannerProfile.MOBILE_2GB, + StorageCapabilities(supportsMappedFiles = false), + ) + // 576 KiB packed is over the 256 KiB off-heap threshold + assertEquals(MemoryDomain.HOST_OFFHEAP, spec.domain) + assertEquals(ScopeKind.MODEL, spec.scope) + } + + @Test + fun heapOnlyPlatformEndsOnTheHeapNoMatterTheSize() { + val spec = AllocationResolver.resolve( + weight(form = WeightForm(residency = WeightResidency.MAPPED)), + PlannerProfile.MOBILE_2GB, + StorageCapabilities.HEAP_ONLY, + ) + assertEquals(MemoryDomain.HOST_HEAP, spec.domain) + } + + @Test + fun dequantizedWeightIsNeverMapped() { + val form = WeightForm(encoding = EncodingRequest.DequantizeTo(FP32), residency = WeightResidency.MAPPED) + val spec = AllocationResolver.resolve(weight(form = form), PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL) + assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain) + // the resident bytes are the dense bytes, and they price the domain decision + assertEquals(Format.dense(FP32), spec.format) + assertEquals(MemoryDomain.HOST_OFFHEAP, spec.domain) // 4 MiB dense is far over threshold + } + + @Test + fun kernelFeedOrderIsALoadTimeCopySoNotMapped() { + val form = WeightForm(order = WeightByteOrder.KERNEL_FEED, residency = WeightResidency.MAPPED) + val spec = AllocationResolver.resolve(weight(form = form), PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL) + assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain) + } + + @Test + fun smallHeapWeightStaysOnTheHeap() { + val dense = Format.dense(FP32) + val bias = PlanTensor("blk.0.attn_q.bias", null, dense, 768, 4L * 768, WeightForm()) + val spec = AllocationResolver.resolve(bias, PlannerProfile.DESKTOP, StorageCapabilities.FULL) + assertEquals(MemoryDomain.HOST_HEAP, spec.domain) + assertEquals(ScopeKind.MODEL, spec.scope) + } + + @Test + fun noFormMeansTheFileBytesByTheProfileRules() { + val spec = AllocationResolver.resolve(weight(form = null), PlannerProfile.DESKTOP, StorageCapabilities.FULL) + assertNotEquals(MemoryDomain.MMAP_FILE, spec.domain) // nothing asked for mapping + assertEquals(MemoryDomain.HOST_OFFHEAP, spec.domain) + } + + @Test + fun transientAllocationsFollowScopeAndThreshold() { + val dense = Format.dense(FP32) + val small = AllocationResolver.resolveTransient(dense, 1024, PlannerProfile.DESKTOP, StorageCapabilities.FULL) + assertEquals(MemoryDomain.HOST_HEAP, small.domain) + assertEquals(ScopeKind.FORWARD, small.scope) + assertTrue(small.mutable) + + val big = AllocationResolver.resolveTransient( + dense, 1L shl 20, PlannerProfile.DESKTOP, StorageCapabilities.FULL, scope = ScopeKind.AMBIENT + ) + assertEquals(MemoryDomain.HOST_OFFHEAP, big.domain) + assertEquals(ScopeKind.AMBIENT, big.scope) + } + + @Test + fun explainSaysWhereAndWhy() { + val mapped = AllocationResolver.explain( + weight(form = WeightForm(residency = WeightResidency.MAPPED)), + PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL, + ) + assertTrue("MMAP_FILE" in mapped && "file's bytes" in mapped, mapped) + + val noMmap = AllocationResolver.explain( + weight(form = WeightForm(residency = WeightResidency.MAPPED)), + PlannerProfile.MOBILE_2GB, StorageCapabilities(supportsMappedFiles = false), + ) + assertTrue("cannot map" in noMmap, noMmap) + + val dequant = AllocationResolver.explain( + weight(form = WeightForm(encoding = EncodingRequest.DequantizeTo(FP32), residency = WeightResidency.MAPPED)), + PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL, + ) + assertTrue("re-encoded at load" in dequant, dequant) + } + + @Test + fun planTensorAllocationDelegatesToTheResolver() { + val w = weight(form = WeightForm(residency = WeightResidency.MAPPED)) + assertEquals( + AllocationResolver.resolve(w, PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL), + w.allocation(PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL), + ) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt index 4630a1748..a0a163857 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/MemoryPlanTest.kt @@ -71,9 +71,14 @@ class MemoryPlanTest { assertEquals(listOf("weights", "kv cache", "forward", "heap"), plan.lines.map { it.section }) assertTrue(plan.lines.first { it.section == "weights" }.resident) assertFalse(plan.lines.first { it.section == "forward" }.resident) - // every weight's allocation is a mapped, model-lifetime, read-only spec - val a = input.weights.first().allocation - assertEquals(MemoryDomain.MMAP_FILE, a.domain); assertEquals(ScopeKind.MODEL, a.scope); assertFalse(a.mutable) + // a weight's allocation is resolved, not assumed (#1143): mapped only when its form asks + // for MAPPED and the platform can map — these weights carry no form, so they fall to the + // profile's heap/off-heap threshold, model-lifetime, read-only + val a = input.weights.first().allocation(PlannerProfile.DESKTOP, StorageCapabilities.FULL) + assertEquals(MemoryDomain.HOST_OFFHEAP, a.domain); assertEquals(ScopeKind.MODEL, a.scope); assertFalse(a.mutable) + val mapped = input.weights.first().copy(form = WeightForm(residency = WeightResidency.MAPPED)) + .allocation(PlannerProfile.MOBILE_2GB, StorageCapabilities.FULL) + assertEquals(MemoryDomain.MMAP_FILE, mapped.domain) } @Test 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 index fa0a58f4f..ba9b44cae 100644 --- 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 @@ -196,4 +196,12 @@ class PlannerProfileTest { assertTrue(text.contains("note: KV cache switched"), text) assertTrue(text.contains("note: weights are counted resident and mapped"), text) } + + @Test + fun domainForSplitsExactlyAtTheOffHeapThreshold() { + val p = PlannerProfile.DESKTOP + assertEquals(sk.ainet.lang.tensor.storage.MemoryDomain.HOST_HEAP, p.domainFor(p.offHeapThresholdBytes - 1)) + assertEquals(sk.ainet.lang.tensor.storage.MemoryDomain.HOST_OFFHEAP, p.domainFor(p.offHeapThresholdBytes)) + assertEquals(sk.ainet.lang.tensor.storage.MemoryDomain.HOST_HEAP, p.domainFor(0)) + } }