From 6b637650064634206f1c18bca547ff7d5d4c4d02 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 10:36:29 +0200 Subject: [PATCH 1/2] =?UTF-8?q?feat(plan):=20AllocationResolver=20?= =?UTF-8?q?=E2=80=94=20one=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)) + } } From 9b35a9ade412db265611b237e8e4fb1be13c4cc0 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 10:50:30 +0200 Subject: [PATCH 2/2] =?UTF-8?q?feat(gguf):=20the=20loader=20consults=20the?= =?UTF-8?q?=20resolved=20allocation=20=E2=80=94=20plan=E2=86=92load=20wiri?= =?UTF-8?q?ng?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plan and load were two disjoint pipelines sharing vocabulary and never talking: nothing on the load path consulted what the resolvers decided. StreamingGgufParametersLoader gains weightFormFor — per-tensor forms with an explicit, documented precedence: your function > the uniform weightForm > the three legacy parameters. The user-wins channel is named as such: whatever you pass outranks every resolver, including 'everything dense on the managed heap'. ResolvedGguf ties it together: header-only planInput, forms resolved from file × profile × kernel capability, overrides applied, and a loader that delivers exactly those forms. The same resolved input is priceable (profiledPlan().requireFits() refuses pre-load) and explainable (explainPlacements() — one line per weight, where and why) before a byte of payload is read. Closes #1144. Co-Authored-By: Claude Fable 5 --- .../sk/ainet/io/gguf/ResolvedGgufLoading.kt | 87 +++++++++ .../io/gguf/StreamingGgufParametersLoader.kt | 78 +++++--- .../gguf/PerTensorFormAndResolvedLoadTest.kt | 178 ++++++++++++++++++ 3 files changed, 317 insertions(+), 26 deletions(-) create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt new file mode 100644 index 000000000..24637a968 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt @@ -0,0 +1,87 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.RandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.AllocationResolver +import sk.ainet.lang.memory.plan.KernelCapabilities +import sk.ainet.lang.memory.plan.PlanInput +import sk.ainet.lang.memory.plan.PlannerProfile +import sk.ainet.lang.memory.plan.ProfiledPlan +import sk.ainet.lang.memory.plan.StorageCapabilities +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.resolveWeightForms +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceSink + +/** + * Plan → load, wired (#1144): the loader consults what the resolvers decided, instead of the two + * pipelines sharing vocabulary and never talking. + * + * [resolve] reads the GGUF *header* (no payload), resolves every weight's [WeightForm] from + * file × [PlannerProfile] × [KernelCapabilities], applies the caller's per-tensor overrides, + * and returns a [Resolution]: a loader that will deliver exactly those forms, plus the resolved + * plan input so the same decisions are priceable ([Resolution.profiledPlan]) and explainable + * ([Resolution.explainPlacements]) *before a byte of payload is read*. + * + * ## Who decides — the precedence order + * + * **Your override > the resolver.** [overrideFormFor] outranks everything for the tensors it names, + * including the deliberately blunt `WeightForm(DequantizeTo(FP32), residency = HEAP)` — everything + * dense, on the managed heap. Tensors it returns `null` for get the resolver's answer. The plan and + * the explanations are computed *after* overrides are applied, so what you print is what you load. + */ +@ExperimentalMemoryApi +public object ResolvedGguf { + + /** What [resolve] decided: the loader that obeys it, and the resolved input that explains it. */ + public data class Resolution( + val loader: StreamingGgufParametersLoader, + /** The header-derived plan input with every weight's resolved (and overridden) form. */ + val input: PlanInput, + val profile: PlannerProfile, + val platform: StorageCapabilities, + ) { + /** The plan these forms cost against [availableBytes] — `requireFits()` refuses pre-load (M2-F6). */ + public fun profiledPlan(availableBytes: Long): ProfiledPlan = profile.plan(input, availableBytes) + + /** One line per weight: where it lands and why — [AllocationResolver.explain] over the resolved forms. */ + public fun explainPlacements(): List = + input.weights.map { AllocationResolver.explain(it, profile, platform) } + } + + /** + * Resolve every weight's form from the header, apply [overrideFormFor], and build the loader + * that delivers it. Reads header only; costs a few kilobytes. + * + * @param capabilities what the backend's kernels can feed — pass the registry-backed + * implementation from the backend in use, or [KernelCapabilities.DENSE_ONLY] to price the + * worst case + * @param overrideFormFor the user-wins channel; `null` per tensor means "resolver decides" + * @throws IllegalStateException when the profile is strict and a weight would dequantize + */ + public fun resolve( + sourceProvider: () -> RandomAccessSource, + profile: PlannerProfile, + capabilities: KernelCapabilities, + ctx: Int? = null, + platform: StorageCapabilities = StorageCapabilities.current(), + overrideFormFor: (tensorName: String) -> WeightForm? = { null }, + onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + traceSink: TraceSink = NoopTraceSink, + ): Resolution { + val resolved = sourceProvider().use { source -> + StreamingGGUFReader.open(source).planInput(ctx) + }.resolveWeightForms(profile, capabilities) + val overridden = resolved.copy( + weights = resolved.weights.map { w -> overrideFormFor(w.name)?.let { w.copy(form = it) } ?: w }, + ) + val forms: Map = overridden.weights.associate { it.name to it.form } + val loader = StreamingGgufParametersLoader( + sourceProvider = sourceProvider, + onProgress = onProgress, + weightFormFor = { name -> forms[name] }, + traceSink = traceSink, + ) + return Resolution(loader, overridden, profile, platform) + } +} diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index f9f3dabbd..7ec59daf9 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -121,6 +121,20 @@ public class StreamingGgufParametersLoader( * silently resolved, so nobody loses a setting they thought they had. */ private val weightForm: WeightForm? = null, + /** + * Per-tensor forms — the *user wins* channel (#1144). + * + * Precedence, explicit and documented: **this function > [weightForm] > the three legacy + * parameters > nothing**. Whatever you return for a tensor outranks every resolver and every + * profile — including the deliberately blunt `WeightForm(DequantizeTo(FP32), residency = HEAP)` + * ("everything dense, on the managed heap"). Return `null` for tensors you have no opinion on; + * they fall through to [weightForm] (or the legacy axes). + * + * The intended producer is `WeightFormResolver`/`resolveWeightForms` via [ResolvedGguf], which + * resolves per tensor from the file × profile × kernel capability — but the contract is the + * same for a hand-written lambda: the loader carries and obeys, it does not decide. + */ + private val weightFormFor: ((tensorName: String) -> WeightForm?)? = null, /** * Where conversions are reported (#1117). * @@ -184,14 +198,15 @@ public class StreamingGgufParametersLoader( }, ) - /** [QuantPolicy.DEQUANTIZE_TO_FP32] asked for through [form], whichever way it was set. */ - private val dequantizeToDense: Boolean = form.encoding is EncodingRequest.DequantizeTo - - /** [StagingPolicy.MAPPED] asked for through [form]. */ - private val mapsTheFile: Boolean = form.residency == WeightResidency.MAPPED - - /** [WeightOrientation.OUT_IN] asked for through [form]. */ - private val reversesWeightShape: Boolean = form.shape == WeightShapeOrientation.OUT_IN + /** + * The form [tensorName] loads under — the precedence order of [weightFormFor], validated the + * same way the uniform [form] was at construction. + */ + private fun formFor(tensorName: String): WeightForm { + val perTensor = weightFormFor?.invoke(tensorName) ?: return form + validateForm(perTensor, "weightFormFor('$tensorName')") + return perTensor + } init { require(quantPolicy != QuantPolicy.RAW_BYTES) { @@ -199,32 +214,36 @@ public class StreamingGgufParametersLoader( "tensors are preserved as packed block TensorData (NATIVE_OPTIMIZED) or " + "dequantized to dense FP32 (DEQUANTIZE_TO_FP32)." } - if (weightForm != null) { + if (weightForm != null || weightFormFor != null) { require( quantPolicy == QuantPolicy.NATIVE_OPTIMIZED && staging == StagingPolicy.HEAP && weightOrientation == WeightOrientation.AS_STORED, ) { - "a WeightForm and the quantPolicy/staging/weightOrientation parameters were both " + - "set. They are the same three axes, so one of them would have been silently " + - "ignored; pass only the form." + "a WeightForm (or weightFormFor) and the quantPolicy/staging/weightOrientation " + + "parameters were both set. They are the same three axes, so one of them would " + + "have been silently ignored; pass only the form." } } + validateForm(form, "weightForm") + } + + private fun validateForm(form: WeightForm, where: String) { require(form.order == WeightByteOrder.AS_STORED || form.shape == WeightShapeOrientation.OUT_IN) { - "WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is defined " + - "relative to a [out, in] weight — which block is 'block b of output row o' has no " + - "answer while the tensor is still labelled in the file's `ne` order." + "$where: WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is " + + "defined relative to a [out, in] weight — which block is 'block b of output row o' " + + "has no answer while the tensor is still labelled in the file's `ne` order." } val requested = form.encoding require(requested !is EncodingRequest.RequantizeTo) { - "EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " + + "$where: EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " + "is not supported by this loader: re-quantizing a weight the file does not already " + "carry needs a quantizer per target encoding, and none of them exist here yet." } val dequantTarget = (requested as? EncodingRequest.DequantizeTo)?.dtype require(dequantTarget == null || dequantTarget == FP32) { - "EncodingRequest.DequantizeTo(${dequantTarget?.name}) is not supported: this loader " + - "dequantizes to FP32 only." + "$where: EncodingRequest.DequantizeTo(${dequantTarget?.name}) is not supported: this " + + "loader dequantizes to FP32 only." } } @@ -233,9 +252,10 @@ public class StreamingGgufParametersLoader( * is reversed for `OUT_IN`, everything else is passed through as the file has it. Only 2-D * tensors are touched — a 1-D bias or norm has no orientation to get wrong. */ - private fun shapeOf(tensorInfo: StreamingTensorInfo): Shape { + private fun shapeOf(tensorInfo: StreamingTensorInfo, form: WeightForm): Shape { val dims = tensorInfo.shape.map { it.toInt() } - val ordered = if (reversesWeightShape && dims.size == 2) dims.reversed() else dims + val reverses = form.shape == WeightShapeOrientation.OUT_IN + val ordered = if (reverses && dims.size == 2) dims.reversed() else dims return Shape(*ordered.toIntArray()) } @@ -248,7 +268,9 @@ public class StreamingGgufParametersLoader( val source = sourceProvider() // MAPPED staging needs a file to map; a Blob or an in-memory source has no path and // silently stays on the heap, which is the documented fallback rather than a failure. - val mapped = if (mapsTheFile) source.filePath?.let { openMappedFile(it) } else null + // With per-tensor forms the file is mapped whenever any tensor *might* ask for it. + val mightMap = form.residency == WeightResidency.MAPPED || weightFormFor != null + val mapped = if (mightMap) source.filePath?.let { openMappedFile(it) } else null try { StreamingGGUFReader.open(source).use { reader -> val tensors = reader.tensors @@ -257,12 +279,15 @@ public class StreamingGgufParametersLoader( var current = 0L for (tensorInfo in tensors) { - val shape = shapeOf(tensorInfo) + val tensorForm = formFor(tensorInfo.name) + val shape = shapeOf(tensorInfo, tensorForm) // A dense F32 tensor under MAPPED staging never reaches the heap: it is a view over // file-backed pages. Everything else reads its bytes (out of the mapping when there // is one — one page-cache copy instead of a channel read). val mappedFloats: Tensor? = - if (mapped != null && tensorInfo.tensorType == GGMLQuantizationType.F32 && dtype == FP32::class) { + if (mapped != null && tensorForm.residency == WeightResidency.MAPPED && + tensorInfo.tensorType == GGMLQuantizationType.F32 && dtype == FP32::class + ) { @Suppress("UNCHECKED_CAST") ctx.fromData( mapped.denseFloats(tensorInfo.absoluteDataOffset, shape) as sk.ainet.lang.tensor.data.TensorData, @@ -340,7 +365,7 @@ public class StreamingGgufParametersLoader( GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q5_1, GGMLQuantizationType.TQ1_0, - GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes) + GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes, tensorForm) else -> throw IllegalStateException( "StreamingGgufParametersLoader: tensor '${tensorInfo.name}' of type " + @@ -382,8 +407,9 @@ public class StreamingGgufParametersLoader( shape: Shape, tensorInfo: StreamingTensorInfo, rawBytes: ByteArray, + tensorForm: WeightForm, ): Tensor { - if (dequantizeToDense && + if (tensorForm.encoding is EncodingRequest.DequantizeTo && (dtype == FP32::class || dtype == FP16::class) ) { val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) @@ -431,7 +457,7 @@ public class StreamingGgufParametersLoader( "quantizedTensor called for non-quantized type ${tensorInfo.tensorType}" ) } - val delivered = if (form.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed + val delivered = if (tensorForm.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed return ctx.fromData(delivered as sk.ainet.lang.tensor.data.TensorData, dtype) } diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt new file mode 100644 index 000000000..a72adecb6 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt @@ -0,0 +1,178 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.KernelCapabilities +import sk.ainet.lang.memory.plan.PlannerProfile +import sk.ainet.lang.memory.plan.StorageCapabilities +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightResidency +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1144: plan → load wiring. A per-tensor form function outranks the uniform form; `ResolvedGguf` + * feeds the loader what the resolvers decided; the user's override outranks the resolver; and the + * decisions are priceable and explainable before the payload is read. + */ +@OptIn(ExperimentalMemoryApi::class) +class PerTensorFormAndResolvedLoadTest { + + private fun file(): File = SyntheticGguf.write( + SyntheticGguf.tensor("w_f32", GGMLQuantizationType.F32, elements = 1024), + SyntheticGguf.tensor("w_q4k", GGMLQuantizationType.Q4_K, elements = 1024), + SyntheticGguf.tensor("w_q80", GGMLQuantizationType.Q8_0, elements = 768) + .copy(dims = listOf(256L, 3L)), + ) + + private fun loadVia(f: File, build: (() -> JvmRandomAccessSource) -> StreamingGgufParametersLoader): + Map> { + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + build { JvmRandomAccessSource.open(f) } + .load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + @Test + fun `a uniform weightFormFor is bit-identical to the single form`() { + val f = file() + try { + for (form in listOf( + WeightForm.AS_STORED_ON_HEAP, + WeightForm(residency = WeightResidency.MAPPED), + WeightForm(encoding = EncodingRequest.DequantizeTo(FP32)), + )) { + val uniform = loadVia(f) { src -> + StreamingGgufParametersLoader(sourceProvider = src, weightForm = form) + } + val perTensor = loadVia(f) { src -> + StreamingGgufParametersLoader(sourceProvider = src, weightFormFor = { form }) + } + assertEquals(uniform.keys, perTensor.keys, "$form: different tensors came out") + for ((name, u) in uniform) { + assertContentEquals( + u.data.copyToFloatArray(), + perTensor.getValue(name).data.copyToFloatArray(), + "$form: $name values", + ) + assertEquals(u.shape, perTensor.getValue(name).shape, "$form: $name shape") + } + } + } finally { + f.delete() + } + } + + @Test + fun `per-tensor forms are honoured per tensor`() { + val f = file() + try { + val loaded = loadVia(f) { src -> + StreamingGgufParametersLoader( + sourceProvider = src, + weightFormFor = { name -> + when (name) { + "w_q4k" -> WeightForm(encoding = EncodingRequest.DequantizeTo(FP32)) + else -> null // uniform default: as stored, on heap + } + }, + ) + } + assertTrue(loaded.getValue("w_q4k").data !is PackedBlockStorage, "w_q4k should be dense") + assertTrue(loaded.getValue("w_q80").data is PackedBlockStorage, "w_q80 should stay packed") + } finally { + f.delete() + } + } + + @Test + fun `ResolvedGguf resolves and the loader obeys`() { + val f = file() + try { + // DENSE_ONLY kernels: the resolver must dequantize every packed weight + val resolution = ResolvedGguf.resolve( + sourceProvider = { JvmRandomAccessSource.open(f) }, + profile = PlannerProfile.DESKTOP, + capabilities = KernelCapabilities.DENSE_ONLY, + platform = StorageCapabilities.FULL, + ) + val q4k = resolution.input.weights.first { it.name == "w_q4k" } + assertTrue(q4k.form?.encoding is EncodingRequest.DequantizeTo, "resolver should dequantize q4k") + assertTrue(q4k.residentBytes > q4k.bytes, "dense costs more than packed and the plan must say so") + + val loaded = LinkedHashMap>() + runBlocking { + resolution.loader.load(DefaultDataExecutionContext(), FP32::class) { name, tensor -> + loaded[name] = tensor + } + } + assertTrue(loaded.getValue("w_q4k").data !is PackedBlockStorage, "loader must obey the resolved form") + + val explains = resolution.explainPlacements() + assertEquals(resolution.input.weights.size, explains.size) + assertTrue(explains.any { "w_q4k" in it }, explains.joinToString("\n")) + + val plan = resolution.profiledPlan(availableBytes = 8L * 1024 * 1024 * 1024) + assertEquals(resolution.input.weights.sumOf { it.residentBytes }, plan.plan.weightsBytes) + } finally { + f.delete() + } + } + + @Test + fun `the override outranks the resolver`() { + val f = file() + try { + val resolution = ResolvedGguf.resolve( + sourceProvider = { JvmRandomAccessSource.open(f) }, + profile = PlannerProfile.DESKTOP, + capabilities = KernelCapabilities.DENSE_ONLY, + platform = StorageCapabilities.FULL, + overrideFormFor = { name -> if (name == "w_q4k") WeightForm.AS_STORED_ON_HEAP else null }, + ) + val q4k = resolution.input.weights.first { it.name == "w_q4k" } + assertEquals(EncodingRequest.KeepAsStored, q4k.form?.encoding, "override must win") + + val loaded = LinkedHashMap>() + runBlocking { + resolution.loader.load(DefaultDataExecutionContext(), FP32::class) { name, tensor -> + loaded[name] = tensor + } + } + assertTrue(loaded.getValue("w_q4k").data is PackedBlockStorage, "override said keep packed") + assertTrue(loaded.getValue("w_q80").data !is PackedBlockStorage, "un-overridden tensors follow the resolver") + } finally { + f.delete() + } + } + + @Test + fun `a strict profile refuses at resolve time — before any payload is read`() { + val f = file() + try { + assertFailsWith { + ResolvedGguf.resolve( + sourceProvider = { JvmRandomAccessSource.open(f) }, + profile = PlannerProfile.MOBILE_2GB, // strict + capabilities = KernelCapabilities.DENSE_ONLY, + platform = StorageCapabilities.FULL, + ) + } + } finally { + f.delete() + } + } +}