From 60950bd93a4817b0f6d193ce2e8f0df2015bc0d9 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 25 Aug 2026 12:51:10 +0200 Subject: [PATCH] feat(memory): WeightForm and a resolver that asks the target, not the caller MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1114. Slice 1 of 5 for #1109. How a weight ends up in memory was decided by three independent flags whoever constructed the loader happened to set: QuantPolicy for dequantization, StagingPolicy for heap vs mapped, WeightOrientation for the shape. Nothing asked the device. A format with no kernel became a per-call dequantization nobody declared; a format with a good kernel could still be handed bytes in an order it does not read. And WeightOrientation only reverses dimensions — the packed byte order that decides whether a kernel reads the right blocks (#973) had no name in that vocabulary at all. WeightForm names all three together, and WeightFormResolver decides them from what the file holds, what PlannerProfile says about the device, and what the backend's kernels can feed. Residency follows the device alone. A dense weight asks for nothing. An encoding with a kernel keeps its encoding and gets KERNEL_FEED order, so the per-weight relayout (#1096) has nothing to do on first use. An encoding with no kernel is dequantized once at load instead of once per forward pass — unless the profile is strict, which finally gives PlannerProfile.strict something to do at load time; nothing consulted it there before. Placement follows the dependency graph rather than the issue text: io-core and backend-api are siblings over lang-core, and the resolver needs PlannerProfile and kernel capabilities while its output is consumed by the loader. lang-core is the only module all three can see. KernelCapabilities is declared there; the registry-backed implementation is in backend-api, where the providers are. That implementation consults both registries, because there are two and a test caught it consulting one. KernelRegistry holds the provider SPI; KernelDispatch holds KernelKey-addressed view kernels, and is where KernelPacks.installReference puts the reference FP32 GEMM. Asking only the first reports a target carrying the reference pack as unable to multiply dense floats — false, and exactly the drift the object exists to avoid. Inert by design: nothing calls the resolver yet. The loader is #1115. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- .../api/kernel/RegistryKernelCapabilities.kt | 52 ++++++++ .../kernel/RegistryKernelCapabilitiesTest.kt | 76 +++++++++++ .../api/jvm/skainet-lang-core.api | 91 +++++++++++++ .../lang/memory/plan/KernelCapabilities.kt | 52 ++++++++ .../sk/ainet/lang/memory/plan/WeightForm.kt | 122 ++++++++++++++++++ .../lang/memory/plan/WeightFormResolver.kt | 74 +++++++++++ .../memory/plan/WeightFormResolverTest.kt | 122 ++++++++++++++++++ 7 files changed, 589 insertions(+) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilities.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilitiesTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/KernelCapabilities.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightForm.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/WeightFormResolverTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilities.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilities.kt new file mode 100644 index 000000000..18c8afdf7 --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilities.kt @@ -0,0 +1,52 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.KernelCapabilities +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * [KernelCapabilities] answered by the kernels actually registered on this device (#1109). + * + * Deliberately thin. `KernelProvider.supports("matmul", [input, weight])` already exists and is + * already the contract providers override when they ship kernels beyond the built-in accessors — + * ternary packs among them — so asking it is asking the same source dispatch will ask. Inventing a + * second capability table beside that one is how the two drift apart. + * + * ## Both registries, because there are two + * + * A matmul kernel can reach this device by either of two routes, and a capability answer that knows + * about one of them is wrong about the other. [KernelRegistry] holds the provider SPI — the + * per-encoding accessors and `supports` — and is what the eager quantized paths consult. + * [KernelDispatch] holds `KernelKey`-addressed view kernels, and is where `KernelPacks.installReference` + * puts the reference FP32 GEMM that every target is supposed to have. Ask only the first and a + * target carrying nothing but the reference pack is reported as unable to multiply dense floats, + * which is both false and exactly the kind of drift this object exists to avoid. + * + * Availability matters as much as registration: a provider whose `isAvailable()` is false on this + * CPU cannot feed anything, whatever it declares. Dispatch kernels carry their requirements in + * their key's capability set instead, and are filtered when they are selected. + */ +@ExperimentalMemoryApi +public object RegistryKernelCapabilities : KernelCapabilities { + + /** The activation dtype every packed matmul kernel in the tree takes. */ + private const val FP32_KEY: String = "Float32" + + override fun canFeedMatmul(encoding: TensorEncoding?): Boolean { + val weightKey = when (encoding) { + null -> FP32_KEY + is TensorEncoding.Dense -> FP32_KEY + else -> encoding.name + } + val fromProviders = KernelRegistry.providers().any { provider -> + provider.isAvailable() && provider.supports("matmul", listOf(FP32_KEY, weightKey)) + } + return fromProviders || dispatchHasMatmulFor(weightKey) + } + + /** Whether a registered [ViewKernel] computes a matmul whose weight operand is [weightKey]. */ + private fun dispatchHasMatmulFor(weightKey: String): Boolean = + KernelDispatch.kernels().any { kernel -> + kernel.key.op == "matmul" && kernel.key.operands.any { it.format.kernelEncodingName == weightKey } + } +} diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilitiesTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilitiesTest.kt new file mode 100644 index 000000000..97f05f99b --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/RegistryKernelCapabilitiesTest.kt @@ -0,0 +1,76 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.storage.TensorEncoding +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * #1109 slice 1: the capability answer comes from the registry, so it cannot disagree with dispatch. + */ +@OptIn(ExperimentalMemoryApi::class) +class RegistryKernelCapabilitiesTest { + + @AfterTest fun cleanup() { KernelDispatch.clearForTesting(); KernelRegistry.clearForTesting() } + + /** A provider that declares one packed matmul and, optionally, is not available here. */ + private class OneEncodingProvider( + private val weightKey: String, + private val available: Boolean = true, + ) : KernelProvider { + override val name: String = "fake-$weightKey" + override val priority: Int = 100 + override fun isAvailable(): Boolean = available + override fun matmulFp32(): Fp32MatmulKernel? = null + override fun supports(opName: String, dtypeKeys: List): Boolean = + opName == "matmul" && dtypeKeys == listOf("Float32", weightKey) + } + + @Test + fun anEncodingWithARegisteredKernelCanBeFed() { + KernelRegistry.clearForTesting() + KernelRegistry.register(OneEncodingProvider("Q4_K")) + + assertTrue(RegistryKernelCapabilities.canFeedMatmul(TensorEncoding.Q4_K)) + assertFalse( + RegistryKernelCapabilities.canFeedMatmul(TensorEncoding.Q6_K), + "nothing registered a Q6_K kernel, so the honest answer is no", + ) + } + + @Test + fun aProviderThatIsNotAvailableHereCannotFeedAnything() { + // Registration is a claim; availability is the fact. A pack compiled for a CPU feature this + // device lacks must not make the resolver keep an encoding nothing can compute (§5.2, #920). + KernelRegistry.clearForTesting() + KernelRegistry.register(OneEncodingProvider("Q4_K", available = false)) + + assertFalse(RegistryKernelCapabilities.canFeedMatmul(TensorEncoding.Q4_K)) + } + + @Test + fun aFeedableBlockedEncodingWantsKernelFeedOrder() { + KernelRegistry.clearForTesting() + KernelRegistry.register(OneEncodingProvider("Q8_0")) + + assertTrue( + RegistryKernelCapabilities.wantsKernelFeedOrder(TensorEncoding.Q8_0), + "every packed matmul kernel in the tree reads input-block-major bytes (#973)", + ) + assertFalse( + RegistryKernelCapabilities.wantsKernelFeedOrder(TensorEncoding.Q4_K), + "an encoding it cannot feed has no order preference to state", + ) + } + + @Test + fun denseIsAnsweredByTheFp32Kernel() { + KernelRegistry.clearForTesting() + assertFalse(RegistryKernelCapabilities.canFeedMatmul(null), "an empty registry can feed nothing") + + KernelPacks.installReference() + assertTrue(RegistryKernelCapabilities.canFeedMatmul(null), "the reference pack always carries FP32") + } +} 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 bf91e9b0f..0c2a45f2d 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 @@ -1471,6 +1471,53 @@ public final class sk/ainet/lang/memory/plan/DeviceMemory { public final class sk/ainet/lang/memory/plan/DeviceMemory$Companion { } +public abstract interface class sk/ainet/lang/memory/plan/EncodingRequest { +} + +public final class sk/ainet/lang/memory/plan/EncodingRequest$DequantizeTo : sk/ainet/lang/memory/plan/EncodingRequest { + public fun (Lsk/ainet/lang/types/DType;)V + public final fun component1 ()Lsk/ainet/lang/types/DType; + public final fun copy (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/memory/plan/EncodingRequest$DequantizeTo; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/EncodingRequest$DequantizeTo;Lsk/ainet/lang/types/DType;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/EncodingRequest$DequantizeTo; + public fun equals (Ljava/lang/Object;)Z + public final fun getDtype ()Lsk/ainet/lang/types/DType; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/EncodingRequest$KeepAsStored : sk/ainet/lang/memory/plan/EncodingRequest { + public static final field INSTANCE Lsk/ainet/lang/memory/plan/EncodingRequest$KeepAsStored; + public fun equals (Ljava/lang/Object;)Z + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/EncodingRequest$RequantizeTo : sk/ainet/lang/memory/plan/EncodingRequest { + public fun (Lsk/ainet/lang/tensor/storage/TensorEncoding;)V + public final fun component1 ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public final fun copy (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Lsk/ainet/lang/memory/plan/EncodingRequest$RequantizeTo; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/EncodingRequest$RequantizeTo;Lsk/ainet/lang/tensor/storage/TensorEncoding;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/EncodingRequest$RequantizeTo; + public fun equals (Ljava/lang/Object;)Z + public final fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + +public abstract interface class sk/ainet/lang/memory/plan/KernelCapabilities { + public static final field Companion Lsk/ainet/lang/memory/plan/KernelCapabilities$Companion; + public abstract fun canFeedMatmul (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Z + public fun wantsKernelFeedOrder (Lsk/ainet/lang/tensor/storage/TensorEncoding;)Z +} + +public final class sk/ainet/lang/memory/plan/KernelCapabilities$Companion { + public final fun getDENSE_ONLY ()Lsk/ainet/lang/memory/plan/KernelCapabilities; + public final fun getEVERYTHING ()Lsk/ainet/lang/memory/plan/KernelCapabilities; +} + +public final class sk/ainet/lang/memory/plan/KernelCapabilities$DefaultImpls { + public static fun wantsKernelFeedOrder (Lsk/ainet/lang/memory/plan/KernelCapabilities;Lsk/ainet/lang/tensor/storage/TensorEncoding;)Z +} + 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; @@ -1758,6 +1805,50 @@ public final class sk/ainet/lang/memory/plan/Suggestion { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/memory/plan/WeightByteOrder : java/lang/Enum { + public static final field AS_STORED Lsk/ainet/lang/memory/plan/WeightByteOrder; + public static final field KERNEL_FEED Lsk/ainet/lang/memory/plan/WeightByteOrder; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/memory/plan/WeightByteOrder; + public static fun values ()[Lsk/ainet/lang/memory/plan/WeightByteOrder; +} + +public final class sk/ainet/lang/memory/plan/WeightForm { + public static final field Companion Lsk/ainet/lang/memory/plan/WeightForm$Companion; + public fun ()V + public fun (Lsk/ainet/lang/memory/plan/EncodingRequest;Lsk/ainet/lang/memory/plan/WeightByteOrder;Lsk/ainet/lang/memory/plan/WeightResidency;)V + public synthetic fun (Lsk/ainet/lang/memory/plan/EncodingRequest;Lsk/ainet/lang/memory/plan/WeightByteOrder;Lsk/ainet/lang/memory/plan/WeightResidency;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()Lsk/ainet/lang/memory/plan/EncodingRequest; + public final fun component2 ()Lsk/ainet/lang/memory/plan/WeightByteOrder; + public final fun component3 ()Lsk/ainet/lang/memory/plan/WeightResidency; + public final fun copy (Lsk/ainet/lang/memory/plan/EncodingRequest;Lsk/ainet/lang/memory/plan/WeightByteOrder;Lsk/ainet/lang/memory/plan/WeightResidency;)Lsk/ainet/lang/memory/plan/WeightForm; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/WeightForm;Lsk/ainet/lang/memory/plan/EncodingRequest;Lsk/ainet/lang/memory/plan/WeightByteOrder;Lsk/ainet/lang/memory/plan/WeightResidency;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/WeightForm; + public fun equals (Ljava/lang/Object;)Z + public final fun getEncoding ()Lsk/ainet/lang/memory/plan/EncodingRequest; + public final fun getOrder ()Lsk/ainet/lang/memory/plan/WeightByteOrder; + public final fun getResidency ()Lsk/ainet/lang/memory/plan/WeightResidency; + public fun hashCode ()I + public final fun isPassThrough ()Z + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/plan/WeightForm$Companion { + public final fun getAS_STORED_ON_HEAP ()Lsk/ainet/lang/memory/plan/WeightForm; +} + +public final class sk/ainet/lang/memory/plan/WeightFormResolver { + public static final field INSTANCE Lsk/ainet/lang/memory/plan/WeightFormResolver; + public final fun resolve (Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/plan/PlannerProfile;Lsk/ainet/lang/memory/plan/KernelCapabilities;)Lsk/ainet/lang/memory/plan/WeightForm; +} + +public final class sk/ainet/lang/memory/plan/WeightResidency : java/lang/Enum { + public static final field HEAP Lsk/ainet/lang/memory/plan/WeightResidency; + public static final field MAPPED Lsk/ainet/lang/memory/plan/WeightResidency; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/memory/plan/WeightResidency; + public static fun values ()[Lsk/ainet/lang/memory/plan/WeightResidency; +} + public final class sk/ainet/lang/memory/trace/CompositeTraceSink : sk/ainet/lang/memory/trace/TraceSink { public fun (Ljava/util/List;)V public fun ([Lsk/ainet/lang/memory/trace/TraceSink;)V diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/KernelCapabilities.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/KernelCapabilities.kt new file mode 100644 index 000000000..f8ac40885 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/KernelCapabilities.kt @@ -0,0 +1,52 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.blockSpec +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * What a backend's kernels can actually be fed (#1109). + * + * The question [WeightFormResolver] needs and nothing could answer before: *which encodings can + * this target compute a matmul on without help?* A backend that has no kernel for the format on + * disk will still produce correct output — by dequantizing, on every forward pass — so the absence + * is invisible until someone profiles it. Asking first turns that into a decision made once. + * + * Declared here rather than in `skainet-backend-api` because the resolver lives beside + * [PlannerProfile] and the two are used together, while the backend modules are downstream. The + * registry-backed implementation is over there, where the providers are. + */ +@ExperimentalMemoryApi +public interface KernelCapabilities { + + /** + * Can this target feed a matmul on a weight encoded as [encoding], with FP32 activations? + * + * `null` means dense — the question is then whether an FP32 matmul kernel exists at all, which + * for any real backend it does. + */ + public fun canFeedMatmul(encoding: TensorEncoding?): Boolean + + /** + * Does the kernel for [encoding] read input-block-major bytes? + * + * Every packed matmul kernel in the tree does (#973), so the default is "yes, if it is blocked + * and we can feed it". A backend whose packed kernels read canonical order overrides this and + * gets `AS_STORED` bytes instead of a pointless permutation. + */ + public fun wantsKernelFeedOrder(encoding: TensorEncoding): Boolean = + encoding.blockSpec != null && canFeedMatmul(encoding) + + public companion object { + + /** A target with dense kernels and nothing else — the conservative assumption. */ + public val DENSE_ONLY: KernelCapabilities = object : KernelCapabilities { + override fun canFeedMatmul(encoding: TensorEncoding?): Boolean = encoding == null + } + + /** A target that can feed anything. Useful in tests; true of no real backend. */ + public val EVERYTHING: KernelCapabilities = object : KernelCapabilities { + override fun canFeedMatmul(encoding: TensorEncoding?): Boolean = true + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightForm.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightForm.kt new file mode 100644 index 000000000..bcf4081b1 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightForm.kt @@ -0,0 +1,122 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.DType + +/** + * The form a weight is asked to take in memory — what it is encoded as, what order its bytes are + * in, and where they live (#1109). + * + * ## Why this is one type and not three flags + * + * The three decisions were already being made, separately, by whoever constructed the loader: + * `QuantPolicy` said whether to dequantize, `StagingPolicy` said heap or mapped, and + * `WeightOrientation` said which way round the shape was. Three problems with that. + * + * Nobody asked the *target*. A weight arrives in whatever the file holds, and nothing consults the + * backend about which encodings its kernels can actually feed — so a format with no kernel becomes + * a per-call dequantization that nobody declared, and a format with a good kernel can still be + * handed bytes in the wrong order. + * + * `WeightOrientation` also stops at the shape: it reverses a 2-D weight's dimensions and leaves the + * bytes alone. The order that decides whether a packed kernel reads the right blocks (#973) had no + * name in that vocabulary at all. [order] gives it one. + * + * And it was the caller's decision when it should have been a resolved one. What the file holds, + * what the device has and what the plan can afford are all knowable at load time, in one place — + * asking the user to pick from three enums is asking them to do that resolution by hand, for a + * device they may not be building for. [WeightFormResolver] does it instead. + * + * ## Repacking and dequantizing are allowed + * + * Both are legitimate; a dequantization is sometimes exactly right. The point is that they become + * *stated intent*, executed once at load, rather than accidents discovered per forward pass — and + * that the cost is visible before it is paid, because a dequantization can quadruple a tensor. + * + * @property encoding what the bytes should encode once loaded + * @property order which way the packed blocks run + * @property residency whether the bytes live on the heap or in file-backed pages + */ +@ExperimentalMemoryApi +public data class WeightForm( + val encoding: EncodingRequest = EncodingRequest.KeepAsStored, + val order: WeightByteOrder = WeightByteOrder.AS_STORED, + val residency: WeightResidency = WeightResidency.HEAP, +) { + /** True when this form asks for nothing — the bytes are used exactly as the file holds them. */ + public val isPassThrough: Boolean + get() = encoding == EncodingRequest.KeepAsStored && + order == WeightByteOrder.AS_STORED && + residency == WeightResidency.HEAP + + public companion object { + /** What every loader did before #1109: the file's bytes, its order, on the heap. */ + public val AS_STORED_ON_HEAP: WeightForm = WeightForm() + } +} + +/** What the loaded bytes should encode. */ +@ExperimentalMemoryApi +public sealed interface EncodingRequest { + + /** Whatever the file holds, untouched. The default, and the only one that costs nothing. */ + public data object KeepAsStored : EncodingRequest + + /** + * Decode to a dense [dtype] at load. + * + * Correct when no kernel can feed the stored encoding — the alternative is dequantizing on + * every forward pass instead of once — and expensive in the obvious way: a Q4_K tensor becomes + * roughly eight times its size as FP32. That is why a resolver that chooses this should be able + * to say so before the load rather than after the OOM. + */ + public data class DequantizeTo(val dtype: DType) : EncodingRequest + + /** + * Re-encode to [encoding] at load — a quantization the file did not already have. + * + * Distinct from [DequantizeTo] in direction and from [KeepAsStored] in cost: it re-quantizes, + * so it is lossy, and it is only ever right when the target's kernels want a format the file + * does not carry. + */ + public data class RequantizeTo(val encoding: TensorEncoding) : EncodingRequest +} + +/** + * Which order a packed weight's blocks run in. + * + * This is the distinction `WeightOrientation` could not express. It reverses *dimensions*; this is + * about *bytes*, and for a block-quantized weight the two are independent: canonical storage and + * kernel feed order hold the same blocks in different physical positions, and they coincide only at + * one block per row (#973, #968). + */ +@ExperimentalMemoryApi +public enum class WeightByteOrder { + + /** The file's own order — canonical row-major blocks, as every GGUF-shaped producer writes. */ + AS_STORED, + + /** + * The order the packed matmul kernels read: input-block-major, every output row's block for one + * input block contiguous. + * + * Asking for this at load is what makes the #1096 relayout unnecessary — the weight arrives in + * feed order and the first forward pass has nothing to convert. + */ + KERNEL_FEED, +} + +/** Where a weight's bytes live. The loader-side spelling of `StagingPolicy` (#1037). */ +@ExperimentalMemoryApi +public enum class WeightResidency { + + /** Read onto the managed heap. The historical behaviour, and the only option in a browser. */ + HEAP, + + /** + * Serve from file-backed pages, which the OS pages in on demand and evicts under pressure — + * the difference between fitting a model on a 2 GB device and not (#921, #922). + */ + MAPPED, +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt new file mode 100644 index 000000000..4a7663235 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightFormResolver.kt @@ -0,0 +1,74 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.blockSpec +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 + +/** + * Decides the [WeightForm] a weight should take, from what the file holds, what the device is, and + * what the backend's kernels can feed (#1109). + * + * This is the piece that lets a model author declare nothing. The same `nn { }` graph runs over an + * FP32 checkpoint on a workstation and a Q4_K GGUF on a phone, and the difference in how the + * weights should be held is a function of three things the author does not know and the loader + * does. + */ +@ExperimentalMemoryApi +public object WeightFormResolver { + + /** + * The form a weight stored as [stored] should take on a device described by [profile], given + * what [capabilities] can feed. + * + * The rules, in order: + * + * 1. **Residency** comes from the profile alone. `weightsMapped` is a statement about the + * device — a 2 GB board cannot hold the weights on the heap whatever they are encoded as. + * 2. **A dense weight** is already in the only form it has. + * 3. **A kernel can feed the stored encoding** → keep it. Ask only whether the kernel wants + * its blocks in feed order, which every packed kernel in the tree does, and hand it bytes + * that way so the per-weight relayout (#1096) never has to run. + * 4. **Nothing can feed it** → the backend would otherwise dequantize on *every forward pass*. + * Doing it once at load is strictly better, so that is the default; but it is also a real + * cost — FP32 is roughly eight times a Q4_K tensor — and a profile that says [PlannerProfile.strict] + * means a missing kernel is a bug to surface, not a slow path to take quietly. + * + * @throws IllegalStateException when nothing can feed [stored] and the profile is strict + */ + public fun resolve( + stored: TensorEncoding?, + profile: PlannerProfile, + capabilities: KernelCapabilities, + ): WeightForm { + val residency = if (profile.weightsMapped) WeightResidency.MAPPED else WeightResidency.HEAP + + if (stored == null || stored is TensorEncoding.Dense) { + return WeightForm(EncodingRequest.KeepAsStored, WeightByteOrder.AS_STORED, residency) + } + + if (capabilities.canFeedMatmul(stored)) { + val order = + if (capabilities.wantsKernelFeedOrder(stored)) WeightByteOrder.KERNEL_FEED + else WeightByteOrder.AS_STORED + return WeightForm(EncodingRequest.KeepAsStored, order, residency) + } + + check(!profile.strict) { + "no kernel on this target can feed a ${stored.name} weight, and ${profile.name} is strict: " + + "loading it would dequantize to FP32 — about ${dequantizedTimes(stored)}× the bytes — and the " + + "profile asks to be told rather than to pay that quietly. Register a ${stored.name} kernel, " + + "convert the file, or resolve with a non-strict profile." + } + + // Once, at load, instead of once per forward pass. + return WeightForm(EncodingRequest.DequantizeTo(FP32), WeightByteOrder.AS_STORED, residency) + } + + /** Roughly how much bigger [encoding] gets as dense FP32 — for the message, not for the plan. */ + private fun dequantizedTimes(encoding: TensorEncoding): String { + val bits = encoding.blockSpec?.bitsPerElement ?: return "several" + val tenths = ((32.0 / bits) * 10).toInt() + return "${tenths / 10}.${tenths % 10}" + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/WeightFormResolverTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/WeightFormResolverTest.kt new file mode 100644 index 000000000..1becccebf --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/WeightFormResolverTest.kt @@ -0,0 +1,122 @@ +package sk.ainet.lang.memory.plan + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1109 slice 1: the form is decided from the file, the device and the kernels — not by the caller. + * + * The table below is the specification. Every encoding is resolved with a kernel present and with + * one absent, under both a desktop and a 2 GB profile, and each cell states what should come out + * and why. If a rule changes, this table is where it changes. + */ +@OptIn(ExperimentalMemoryApi::class) +class WeightFormResolverTest { + + private val packed: List = listOf( + TensorEncoding.Q4_0, TensorEncoding.Q5_0, TensorEncoding.Q5_1, TensorEncoding.Q8_0, + TensorEncoding.Q4_K, TensorEncoding.Q5_K, TensorEncoding.Q6_K, + TensorEncoding.TQ1_0, TensorEncoding.TQ2_0, + ) + + /** A target that can feed exactly [supported] and nothing else. */ + private fun capableOf(vararg supported: TensorEncoding): KernelCapabilities = + object : KernelCapabilities { + override fun canFeedMatmul(encoding: TensorEncoding?): Boolean = + encoding == null || encoding in supported + } + + @Test + fun `a weight whose kernel exists keeps its encoding and gets feed order`() { + for (encoding in packed) { + for (profile in listOf(PlannerProfile.DESKTOP, PlannerProfile.MOBILE_2GB)) { + val form = WeightFormResolver.resolve(encoding, profile, capableOf(encoding)) + assertEquals( + EncodingRequest.KeepAsStored, form.encoding, + "${encoding.name} on ${profile.name}: a feedable weight must not be re-encoded", + ) + assertEquals( + WeightByteOrder.KERNEL_FEED, form.order, + "${encoding.name} on ${profile.name}: the kernel reads input-block-major, so load it that way", + ) + } + } + } + + @Test + fun `a weight with no kernel is dequantized once at load rather than every forward pass`() { + for (encoding in packed) { + val form = WeightFormResolver.resolve(encoding, PlannerProfile.DESKTOP, KernelCapabilities.DENSE_ONLY) + assertEquals( + EncodingRequest.DequantizeTo(FP32), form.encoding, + "${encoding.name}: with no kernel the backend would dequantize per call; once is better", + ) + assertEquals(WeightByteOrder.AS_STORED, form.order, "${encoding.name}: dense bytes have no block order") + } + } + + @Test + fun `a strict profile refuses the silent dequantization instead of paying for it`() { + val strict = PlannerProfile.MOBILE_2GB.copy(strict = true) + val failure = assertFailsWith { + WeightFormResolver.resolve(TensorEncoding.Q4_K, strict, KernelCapabilities.DENSE_ONLY) + } + val message = failure.message!! + assertTrue(message.contains("Q4_K"), "it names the encoding: $message") + assertTrue(message.contains("strict"), "and why it refused: $message") + assertTrue(message.contains("×"), "and what it would have cost: $message") + } + + @Test + fun `residency follows the device and nothing else`() { + val mapped = PlannerProfile.DESKTOP.copy(weightsMapped = true) + for (encoding in packed + listOf(null)) { + assertEquals( + WeightResidency.MAPPED, + WeightFormResolver.resolve(encoding, mapped, KernelCapabilities.EVERYTHING).residency, + "${encoding?.name ?: "dense"}: weightsMapped is a statement about the device, not the encoding", + ) + assertEquals( + WeightResidency.HEAP, + WeightFormResolver.resolve(encoding, PlannerProfile.DESKTOP, KernelCapabilities.EVERYTHING).residency, + ) + } + } + + @Test + fun `a dense weight asks for nothing`() { + val form = WeightFormResolver.resolve(null, PlannerProfile.DESKTOP, KernelCapabilities.DENSE_ONLY) + assertTrue(form.isPassThrough, "a dense weight on a plain profile is already in its only form: $form") + assertEquals(WeightForm.AS_STORED_ON_HEAP, form) + } + + @Test + fun `the same weight resolves differently on two targets`() { + // The point of the whole exercise, in one assertion: identical input, different answers, + // and the model author wrote nothing either way. + val withKernel = WeightFormResolver.resolve( + TensorEncoding.Q4_K, PlannerProfile.MOBILE_2GB.copy(weightsMapped = true), capableOf(TensorEncoding.Q4_K), + ) + val withoutKernel = WeightFormResolver.resolve( + TensorEncoding.Q4_K, PlannerProfile.DESKTOP, KernelCapabilities.DENSE_ONLY, + ) + assertTrue(withKernel != withoutKernel, "same file, different devices, same resolved form — that would be the bug") + assertEquals(EncodingRequest.KeepAsStored, withKernel.encoding) + assertEquals(EncodingRequest.DequantizeTo(FP32), withoutKernel.encoding) + } + + @Test + fun `a backend whose packed kernels read canonical order is not handed a pointless permutation`() { + val canonicalReader = object : KernelCapabilities { + override fun canFeedMatmul(encoding: TensorEncoding?): Boolean = true + override fun wantsKernelFeedOrder(encoding: TensorEncoding): Boolean = false + } + val form = WeightFormResolver.resolve(TensorEncoding.Q8_0, PlannerProfile.DESKTOP, canonicalReader) + assertEquals(WeightByteOrder.AS_STORED, form.order, "relayouting for a kernel that does not want it is waste") + } +}