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