diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/QuantPolicy.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/QuantPolicy.kt index b68581ea0..55f669cb2 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/QuantPolicy.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/QuantPolicy.kt @@ -5,6 +5,13 @@ package sk.ainet.io.model * * Shared across all weight loaders (LLaMA, Gemma, etc.). */ +@Deprecated( + "One of the three axes WeightForm replaces (#1109): this is the encoding axis. No ReplaceWith, " + + "because EncodingRequest is not a drop-in — NATIVE_OPTIMIZED becomes KeepAsStored, " + + "DEQUANTIZE_TO_FP32 becomes DequantizeTo(FP32), and RAW_BYTES has no counterpart because " + + "no loader ever supported it. A wrong ReplaceWith would be worse than none.", + level = DeprecationLevel.WARNING, +) public enum class QuantPolicy { /** Keep quantized payloads as raw bytes (Int8 tensor) with quantized shape. */ RAW_BYTES, diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/StagingPolicy.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/StagingPolicy.kt index ca0f973d5..5df2362ed 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/StagingPolicy.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/StagingPolicy.kt @@ -7,6 +7,15 @@ package sk.ainet.io.model * bytes live*. The two are the axes of one loader — `quantPolicy × staging` — instead of the * separate code paths ("streaming loader" vs "mapped weights helper") they used to be. */ +@Deprecated( + "One of the three axes WeightForm replaces (#1109): this is the residency axis. Same two " + + "values, resolved from the device's PlannerProfile rather than set by the caller.", + ReplaceWith( + "WeightResidency", + "sk.ainet.lang.memory.plan.WeightResidency", + ), + DeprecationLevel.WARNING, +) public enum class StagingPolicy { /** Read tensor bytes onto the heap. The historical behaviour, and the only option in a browser. */ HEAP, diff --git a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/WeightOrientation.kt b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/WeightOrientation.kt index fad8d7188..c13d5f7f9 100644 --- a/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/WeightOrientation.kt +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/WeightOrientation.kt @@ -12,6 +12,15 @@ package sk.ainet.io.model * computes the wrong permutation — or refuses, when `out` is not a multiple of the block size. Both * failures are in the census. */ +@Deprecated( + "One of the three axes WeightForm replaces (#1109): this is the shape axis. Same two values, " + + "asked as part of one resolved decision instead of a separate flag.", + ReplaceWith( + "WeightShapeOrientation", + "sk.ainet.lang.memory.plan.WeightShapeOrientation", + ), + DeprecationLevel.WARNING, +) public enum class WeightOrientation { /** * The file's own order, unreversed — GGUF `ne`, so `[in, out]` for a 2-D weight. What the 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 1e13759dd..387d359dc 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 @@ -4,6 +4,11 @@ import sk.ainet.context.ExecutionContext import sk.ainet.io.ParametersLoader import sk.ainet.io.RandomAccessSource import sk.ainet.io.gguf.dequant.DequantOps +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightResidency +import sk.ainet.lang.memory.plan.WeightShapeOrientation import sk.ainet.io.model.QuantPolicy import sk.ainet.io.model.StagingPolicy import sk.ainet.io.model.WeightOrientation @@ -97,14 +102,82 @@ public class StreamingGgufParametersLoader( * changes what every consumer sees. New code should ask for `OUT_IN`. */ private val weightOrientation: WeightOrientation = WeightOrientation.AS_STORED, + /** + * The form weights should take in memory, as one decision instead of three (#1109, #1115). + * + * [quantPolicy], [staging] and [weightOrientation] are the same three axes asked separately, + * and asked of the *caller* — who has to answer for a device they may not be building for. + * `WeightFormResolver.resolve(stored, profile, capabilities)` answers instead, from what the + * file holds, what the device is, and what the backend's kernels can feed. + * + * `null` (the default) means "use the three parameters", so every existing caller is + * byte-identical. Passing a form while also setting any of the three is rejected rather than + * silently resolved, so nobody loses a setting they thought they had. + */ + private val weightForm: WeightForm? = null, ) : ParametersLoader { + /** The three axes as one value: [weightForm] if given, otherwise what the three parameters say. */ + private val form: WeightForm = weightForm ?: WeightForm( + encoding = when (quantPolicy) { + QuantPolicy.DEQUANTIZE_TO_FP32 -> EncodingRequest.DequantizeTo(FP32) + else -> EncodingRequest.KeepAsStored + }, + order = WeightByteOrder.AS_STORED, + shape = when (weightOrientation) { + WeightOrientation.OUT_IN -> WeightShapeOrientation.OUT_IN + else -> WeightShapeOrientation.AS_STORED + }, + residency = when (staging) { + StagingPolicy.MAPPED -> WeightResidency.MAPPED + else -> WeightResidency.HEAP + }, + ) + + /** [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 + init { require(quantPolicy != QuantPolicy.RAW_BYTES) { "StreamingGgufParametersLoader does not support QuantPolicy.RAW_BYTES — quantized " + "tensors are preserved as packed block TensorData (NATIVE_OPTIMIZED) or " + "dequantized to dense FP32 (DEQUANTIZE_TO_FP32)." } + if (weightForm != 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." + } + } + require(form.order == WeightByteOrder.AS_STORED) { + "WeightByteOrder.KERNEL_FEED is not supported by this loader yet (#1120). The bytes are " + + "easy — TensorView.prepack(INPUT_BLOCK_MAJOR) does the permutation — but packed " + + "TensorData addresses packedData as canonical row-major in getBlockScale, getCode " + + "and dequantizeBlock, so feed-order bytes would decode the wrong elements without " + + "failing (#973, #968). Packed storage has to be able to declare its own order first." + } + val requested = form.encoding + require(requested !is EncodingRequest.RequantizeTo) { + "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." + } } /** @@ -114,7 +187,7 @@ public class StreamingGgufParametersLoader( */ private fun shapeOf(tensorInfo: StreamingTensorInfo): Shape { val dims = tensorInfo.shape.map { it.toInt() } - val ordered = if (weightOrientation == WeightOrientation.OUT_IN && dims.size == 2) dims.reversed() else dims + val ordered = if (reversesWeightShape && dims.size == 2) dims.reversed() else dims return Shape(*ordered.toIntArray()) } @@ -127,7 +200,7 @@ 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 (staging == StagingPolicy.MAPPED) source.filePath?.let { openMappedFile(it) } else null + val mapped = if (mapsTheFile) source.filePath?.let { openMappedFile(it) } else null try { StreamingGGUFReader.open(source).use { reader -> val tensors = reader.tensors @@ -254,7 +327,7 @@ public class StreamingGgufParametersLoader( tensorInfo: StreamingTensorInfo, rawBytes: ByteArray, ): Tensor { - if (quantPolicy == QuantPolicy.DEQUANTIZE_TO_FP32 && + if (dequantizeToDense && (dtype == FP32::class || dtype == FP16::class) ) { val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt new file mode 100644 index 000000000..0a4ac1d08 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt @@ -0,0 +1,179 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.model.QuantPolicy +import sk.ainet.io.model.StagingPolicy +import sk.ainet.io.model.WeightOrientation +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightResidency +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.tensor.Tensor +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 + +/** + * #1115: the loader takes one `WeightForm` where it took three flags, and the change is a change of + * *spelling only*. + * + * That is the claim worth testing, because it is the one that can quietly be false. Every + * combination of the three deprecated parameters is loaded twice — once through them, once through + * the `WeightForm` they map to — and the two must agree on shapes and on every element. If the + * mapping is wrong anywhere, some cell of that product disagrees. + */ +@Suppress("DEPRECATION") +@OptIn(ExperimentalMemoryApi::class) +class WeightFormLoaderParityTest { + + /** Mixed encodings, and a 2-D weight so the shape axis has something to reverse. */ + 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)), + SyntheticGguf.tensor("w_f16", GGMLQuantizationType.F16, elements = 1024), + ) + + 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 `every combination of the three flags loads identically through the form it maps to`() { + val f = file() + try { + for (quant in listOf(QuantPolicy.NATIVE_OPTIMIZED, QuantPolicy.DEQUANTIZE_TO_FP32)) { + for (staging in listOf(StagingPolicy.HEAP, StagingPolicy.MAPPED)) { + for (orientation in listOf(WeightOrientation.AS_STORED, WeightOrientation.OUT_IN)) { + val viaFlags = loadVia(f) { src -> + StreamingGgufParametersLoader( + sourceProvider = src, + quantPolicy = quant, + staging = staging, + weightOrientation = orientation, + ) + } + val viaForm = loadVia(f) { src -> + StreamingGgufParametersLoader( + sourceProvider = src, + weightForm = WeightForm( + encoding = when (quant) { + QuantPolicy.DEQUANTIZE_TO_FP32 -> EncodingRequest.DequantizeTo(FP32) + else -> EncodingRequest.KeepAsStored + }, + shape = when (orientation) { + WeightOrientation.OUT_IN -> WeightShapeOrientation.OUT_IN + else -> WeightShapeOrientation.AS_STORED + }, + residency = when (staging) { + StagingPolicy.MAPPED -> WeightResidency.MAPPED + else -> WeightResidency.HEAP + }, + ), + ) + } + + val label = "$quant/$staging/$orientation" + assertEquals(viaFlags.keys, viaForm.keys, "$label: different tensors came out") + for ((name, flagsTensor) in viaFlags) { + val formTensor = viaForm.getValue(name) + assertEquals(flagsTensor.shape, formTensor.shape, "$label: $name shape") + assertContentEquals( + flagsTensor.data.copyToFloatArray(), + formTensor.data.copyToFloatArray(), + "$label: $name values", + ) + } + } + } + } + } finally { + f.delete() + } + } + + @Test + fun `the default loader is the default form`() { + val f = file() + try { + val implicit = loadVia(f) { src -> StreamingGgufParametersLoader(sourceProvider = src) } + val explicit = loadVia(f) { src -> + StreamingGgufParametersLoader(sourceProvider = src, weightForm = WeightForm.AS_STORED_ON_HEAP) + } + for ((name, tensor) in implicit) { + assertContentEquals( + tensor.data.copyToFloatArray(), explicit.getValue(name).data.copyToFloatArray(), name, + ) + } + } finally { + f.delete() + } + } + + @Test + fun `setting both a form and a flag is refused rather than silently resolved`() { + // One of the two would have to lose, and a caller who set a flag believes it is in effect. + val failure = assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file()) }, + staging = StagingPolicy.MAPPED, + weightForm = WeightForm.AS_STORED_ON_HEAP, + ) + } + assertTrue(failure.message!!.contains("pass only the form"), failure.message!!) + } + + @Test + fun `KERNEL_FEED is refused for the reason it is refused`() { + // Not "unsupported": the bytes are the easy part. The refusal is because packed TensorData + // reads packedData as canonical row-major, so feed-order bytes decode wrong silently (#1120). + val failure = assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file()) }, + weightForm = WeightForm(order = WeightByteOrder.KERNEL_FEED), + ) + } + val message = failure.message!! + assertTrue(message.contains("#1120"), "it names where this is being fixed: $message") + assertTrue(message.contains("canonical row-major"), "and why it cannot be faked: $message") + } + + @Test + fun `an encoding request this loader cannot honour is refused up front`() { + val requantize = assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file()) }, + weightForm = WeightForm( + encoding = EncodingRequest.RequantizeTo(sk.ainet.lang.tensor.storage.TensorEncoding.Q8_0), + ), + ) + } + assertTrue(requantize.message!!.contains("quantizer"), requantize.message!!) + + val toFp16 = assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file()) }, + weightForm = WeightForm( + encoding = EncodingRequest.DequantizeTo(sk.ainet.lang.types.FP16), + ), + ) + } + assertTrue(toFp16.message!!.contains("FP32 only"), toFp16.message!!) + } +} 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 0c2a45f2d..f2d4a2637 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 @@ -1816,17 +1816,19 @@ public final class sk/ainet/lang/memory/plan/WeightByteOrder : java/lang/Enum { 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 fun (Lsk/ainet/lang/memory/plan/EncodingRequest;Lsk/ainet/lang/memory/plan/WeightByteOrder;Lsk/ainet/lang/memory/plan/WeightShapeOrientation;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/WeightShapeOrientation;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 final fun component3 ()Lsk/ainet/lang/memory/plan/WeightShapeOrientation; + public final fun component4 ()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/WeightShapeOrientation;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/WeightShapeOrientation;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 final fun getShape ()Lsk/ainet/lang/memory/plan/WeightShapeOrientation; public fun hashCode ()I public final fun isPassThrough ()Z public fun toString ()Ljava/lang/String; @@ -1849,6 +1851,14 @@ public final class sk/ainet/lang/memory/plan/WeightResidency : java/lang/Enum { public static fun values ()[Lsk/ainet/lang/memory/plan/WeightResidency; } +public final class sk/ainet/lang/memory/plan/WeightShapeOrientation : java/lang/Enum { + public static final field AS_STORED Lsk/ainet/lang/memory/plan/WeightShapeOrientation; + public static final field OUT_IN Lsk/ainet/lang/memory/plan/WeightShapeOrientation; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/memory/plan/WeightShapeOrientation; + public static fun values ()[Lsk/ainet/lang/memory/plan/WeightShapeOrientation; +} + 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/WeightForm.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/WeightForm.kt index bcf4081b1..a2cc5d056 100644 --- 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 @@ -36,18 +36,21 @@ import sk.ainet.lang.types.DType * * @property encoding what the bytes should encode once loaded * @property order which way the packed blocks run + * @property shape which way round the dimensions are labelled * @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 shape: WeightShapeOrientation = WeightShapeOrientation.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 && + shape == WeightShapeOrientation.AS_STORED && residency == WeightResidency.HEAP public companion object { @@ -107,6 +110,27 @@ public enum class WeightByteOrder { KERNEL_FEED, } +/** + * Which way round a 2-D weight's dimensions are labelled. + * + * Orthogonal to [WeightByteOrder], and easy to conflate with it: this is about the *shape*, that is + * about the *bytes*. Reversing the dimensions of a packed weight moves no data at all, while + * changing its block order moves every block and leaves the shape alone. A weight can need either, + * both, or neither. + */ +@ExperimentalMemoryApi +public enum class WeightShapeOrientation { + + /** The file's own order, unreversed — GGUF `ne`, so `[in, out]` for a 2-D weight. */ + AS_STORED, + + /** + * Logical `[out, in]`: the convention the engine, HF checkpoints and every kernel assume. The + * bytes are untouched, because they already are `[out, in]` row-major; only the label changes. + */ + OUT_IN, +} + /** Where a weight's bytes live. The loader-side spelling of `StagingPolicy` (#1037). */ @ExperimentalMemoryApi public enum class WeightResidency { 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 index 4a7663235..f0a9f30ad 100644 --- 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 @@ -23,6 +23,11 @@ public object WeightFormResolver { * * The rules, in order: * + * The shape axis ([WeightShapeOrientation]) is deliberately not resolved here. Which way round + * a weight's dimensions are labelled is a property of the *checkpoint convention*, not of the + * device, and reversing it changes every consumer's idea of a tensor's shape — so it stays the + * caller's explicit choice, defaulting to what the file says. + * * 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. @@ -44,14 +49,14 @@ public object WeightFormResolver { 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) + return WeightForm(EncodingRequest.KeepAsStored, WeightByteOrder.AS_STORED, residency = residency) } if (capabilities.canFeedMatmul(stored)) { val order = if (capabilities.wantsKernelFeedOrder(stored)) WeightByteOrder.KERNEL_FEED else WeightByteOrder.AS_STORED - return WeightForm(EncodingRequest.KeepAsStored, order, residency) + return WeightForm(EncodingRequest.KeepAsStored, order, residency = residency) } check(!profile.strict) { @@ -62,7 +67,7 @@ public object WeightFormResolver { } // Once, at load, instead of once per forward pass. - return WeightForm(EncodingRequest.DequantizeTo(FP32), WeightByteOrder.AS_STORED, residency) + return WeightForm(EncodingRequest.DequantizeTo(FP32), WeightByteOrder.AS_STORED, residency = residency) } /** Roughly how much bigger [encoding] gets as dense FP32 — for the message, not for the plan. */