From d7904cb93046bbfb958ba6c018460e873c207052 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 21:17:37 +0200 Subject: [PATCH] feat(io): weight orientation at the load boundary, and a guard that refuses a transposed label MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1098 (#973 census contradiction #6). GGUF writes dimensions in `ne` order, so a weight the engine calls `[out, in]` arrives from the streaming loader labelled `[in, out]` — while its bytes are already `[out, in]` row-major. Only the label is wrong, and the label is what the block relayout reads: driven by `[in, out]` it permutes the wrong grid, or refuses because `out` is not a multiple of the block size. Both failures are in the census. - `WeightOrientation { AS_STORED, OUT_IN }` on `StreamingGgufParametersLoader`. `OUT_IN` reverses 2-D weights at the boundary and touches nothing else — not the bytes, not 1-D tensors, which have no orientation to get wrong. Defaults to `AS_STORED`, today's behaviour, because reversing shapes changes what every consumer sees. - `PackedWeights.requireOutIn(rows, inputDim, encoding)`, run by `prepackForMatmul`: refuses a weight that looks transposed instead of computing a wrong permutation from it, and names the fix. It is a heuristic and says so — it fires when the *first* dimension is block-aligned and the second is not, which is exactly what `ne` order produces, and stays quiet when both are aligned and it cannot tell. A false refusal would be worse than none. - The normative doc gains the orientation section, including the fact that this repository's two GGUF readers disagree about it today: the legacy `GGUFReader` reverses dimensions, the streaming one does not. `WeightOrientation` is how a caller states which it wants. `SyntheticGguf` can now write multi-dimensional tensors, so the test writes a real `ne = [64, 4]` weight and asserts that `AS_STORED` reports `[64, 4]`, `OUT_IN` reports `[4, 64]`, and the decoded values are identical either way. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- docs/design/memory/packed-weight-layout.md | 30 ++++++- .../ainet/backend/api/kernel/PackedWeights.kt | 31 ++++++++ .../backend/api/kernel/PackedWeightsTest.kt | 28 +++++++ .../sk/ainet/io/model/WeightOrientation.kt | 29 +++++++ .../io/gguf/StreamingGgufParametersLoader.kt | 26 +++++- .../kotlin/sk/ainet/io/gguf/SyntheticGguf.kt | 10 ++- .../sk/ainet/io/gguf/WeightOrientationTest.kt | 79 +++++++++++++++++++ 7 files changed, 227 insertions(+), 6 deletions(-) create mode 100644 skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/WeightOrientation.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightOrientationTest.kt diff --git a/docs/design/memory/packed-weight-layout.md b/docs/design/memory/packed-weight-layout.md index bb7aa0911..4a2e6d853 100644 --- a/docs/design/memory/packed-weight-layout.md +++ b/docs/design/memory/packed-weight-layout.md @@ -76,6 +76,28 @@ green-CI hotfix. Every fixture is three blocks wide on purpose. At one block per row the two orders coincide, and a test built that way passes whichever convention the code holds. +## Orientation at the load boundary + +A packed weight is logically `[out, in]`, and its blocks tile **`in`**. GGUF writes dimensions in +`ne` order — fastest-varying first — so the same weight arrives labelled `[in, out]`, while its +*bytes* are already `[out, in]` row-major. Only the label is wrong, and the label is what the +relayout reads: driven by `[in, out]` it permutes the wrong grid, or refuses because `out` is not a +multiple of the block size. Both failures are in #973's census. + +- `StreamingGgufParametersLoader(weightOrientation = WeightOrientation.OUT_IN)` fixes the label at + the boundary, reversing 2-D weights only. Nothing about the bytes changes. It defaults to + `AS_STORED` — today's behaviour — because reversing shapes changes what every consumer sees; new + code should ask for `OUT_IN`. +- `PackedWeights.requireOutIn(rows, inputDim, encoding)` refuses a weight that looks transposed + instead of computing a wrong permutation from it, and names the fix. `prepackForMatmul` runs it. + The check is a heuristic and says so: it fires when the *first* dimension is block-aligned and the + second is not, which is exactly the shape `ne` order produces, and stays quiet when both are + aligned and it cannot tell. + +Note the two GGUF readers in this repository disagree about this today: the legacy `GGUFReader` +reverses dimensions, the streaming one does not. `WeightOrientation` is how a caller states which it +wants rather than discovering it. + ## Rules 1. **A file's bytes are `ROW_MAJOR`.** Anything loaded from GGUF, produced by a quantizer, or @@ -108,8 +130,10 @@ registry through the ordered key since [#1095](https://github.com/SKaiNET-developers/SKaiNET/issues/1095); `PackedWeights` and `PackedLayoutFixtures` since [#1097](https://github.com/SKaiNET-developers/SKaiNET/issues/1097). +Weight orientation at the load boundary is opt-in since +[#1098](https://github.com/SKaiNET-developers/SKaiNET/issues/1098), with a guard that refuses a +wrongly-labelled weight rather than mis-permuting it. + Still open under #973: the weight-transposing matmul primitive that removes the packed `ops.transpose` entirely ([#1096](https://github.com/SKaiNET-developers/SKaiNET/issues/1096)), and -normalizing weight orientation at the load boundary -([#1098](https://github.com/SKaiNET-developers/SKaiNET/issues/1098)) — until that lands, a -verbatim-loaded GGUF weight's `[in, out]` shape still disagrees with what the relayout assumes. +making `OUT_IN` the default once downstream consumers have moved. diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedWeights.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedWeights.kt index 53cefa47e..4fd8dd754 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedWeights.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedWeights.kt @@ -41,6 +41,7 @@ public object PackedWeights { ): TensorView { require(weight.layout.blocked) { "prepackForMatmul takes a block-packed weight, got ${weight.format}" } require(weight.shape.rank == 2) { "a matmul weight is 2-D [out, in], got ${weight.shape}" } + requireOutIn(weight.shape[0], weight.shape[1], weight.format.encoding) return weight.prepack(BlockOrder.INPUT_BLOCK_MAJOR, scope, sink) } @@ -93,6 +94,36 @@ public object PackedWeights { return out } + /** + * Refuse a weight that looks transposed, instead of computing the wrong permutation from it + * (#973 census contradiction #6; #1098). + * + * A packed weight's blocks tile the **input** dimension, so for a correctly oriented + * `[out, in]` weight `in` is a multiple of the block size. When the *first* dimension is + * block-aligned and the second is not, the tensor is almost certainly `[in, out]` — the shape + * a GGUF's `ne` order produces — and relayouting it would silently permute the wrong grid. + * + * The check is a heuristic and says so: a square weight, or one where both dimensions are + * aligned, passes either way. It catches the case that actually shipped. + */ + public fun requireOutIn(rows: Int, inputDim: Int, encoding: TensorEncoding) { + val spec = encoding.blockSpec ?: return + if (spec.isPerTensor) return + val blockSize = spec.blockSize + val inputAligned = inputDim % blockSize == 0 + val rowsAligned = rows % blockSize == 0 + require(inputAligned || !rowsAligned) { + "weight [$rows, $inputDim] looks like [in, out]: ${encoding.name} tiles the *input* dimension in " + + "blocks of $blockSize, and $inputDim is not a multiple of it while $rows is. A GGUF's ne order " + + "produces exactly this — load with WeightOrientation.OUT_IN, or transpose the label before " + + "relayouting (#973, docs/design/memory/packed-weight-layout.md)." + } + require(inputAligned) { + "weight [$rows, $inputDim] cannot be relayouted: ${encoding.name} needs the input dimension to be a " + + "multiple of $blockSize" + } + } + /** Block geometry of [encoding] — what a converter needs to call [toKernelOrder]. */ public fun blocksPerRow(encoding: TensorEncoding, inputDim: Int): Int { val spec = encoding.blockSpec diff --git a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/PackedWeightsTest.kt b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/PackedWeightsTest.kt index 904a12ce9..8952612c1 100644 --- a/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/PackedWeightsTest.kt +++ b/skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/PackedWeightsTest.kt @@ -106,6 +106,34 @@ class PackedWeightsTest { assertEquals(2, PackedWeights.blocksPerRow(TensorEncoding.Q4_K, 512)) } + + @Test + fun aWeightLabelledTheWrongWayRoundIsRefusedRatherThanMisPermuted() { + // #1098 / #973 census #6: a packed weight's blocks tile the *input* dimension, so + // relayouting an [in, out] label permutes the wrong grid — which is what a GGUF's ne order + // hands you today. + val failure = assertFailsWith { + PackedWeights.requireOutIn(rows = 128, inputDim = 3, encoding = TensorEncoding.Q8_0) + } + assertTrue(failure.message!!.contains("looks like [in, out]"), failure.message!!) + assertTrue(failure.message!!.contains("WeightOrientation.OUT_IN"), "and says how to fix it") + + PackedWeights.requireOutIn(rows = 3, inputDim = 128, encoding = TensorEncoding.Q8_0) + } + + @Test + fun theOrientationGuardStaysQuietWhenItCannotTell() { + // both dimensions block-aligned: ambiguous, and a false refusal would be worse than none + PackedWeights.requireOutIn(rows = 64, inputDim = 128, encoding = TensorEncoding.Q8_0) + PackedWeights.requireOutIn(rows = 128, inputDim = 128, encoding = TensorEncoding.Q8_0) + // an unaligned input dimension cannot be relayouted whichever way round it is + assertFailsWith { + PackedWeights.requireOutIn(rows = 5, inputDim = 7, encoding = TensorEncoding.Q8_0) + } + // a dense encoding has no block grid, so there is nothing to check + PackedWeights.requireOutIn(rows = 5, inputDim = 7, encoding = TensorEncoding.Dense(4)) + } + private fun blockSizeOf(encoding: TensorEncoding): Int = when (encoding) { TensorEncoding.Q4_0, TensorEncoding.Q5_0, TensorEncoding.Q5_1, TensorEncoding.Q8_0 -> 32 else -> 256 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 new file mode 100644 index 000000000..fad8d7188 --- /dev/null +++ b/skainet-io/skainet-io-core/src/commonMain/kotlin/sk/ainet/io/model/WeightOrientation.kt @@ -0,0 +1,29 @@ +package sk.ainet.io.model + +/** + * Which way round a loaded 2-D weight's shape is (#973 census contradiction #6; #1098). + * + * GGUF stores dimensions in `ne` order — fastest-varying first — so a weight the rest of the world + * calls `[out, in]` is written as `ne = [in, out]`. The bytes are the same either way: row-major + * with the input dimension fastest, which *is* `[out, in]` row-major. Only the label differs. + * + * That label matters, because everything downstream of the loader assumes `[out, in]`: the block + * grid of a packed weight is `out × blocksPerRow`, so a relayout driven by an `[in, out]` shape + * computes the wrong permutation — or refuses, when `out` is not a multiple of the block size. Both + * failures are in the census. + */ +public enum class WeightOrientation { + /** + * The file's own order, unreversed — GGUF `ne`, so `[in, out]` for a 2-D weight. What the + * streaming loader has always produced, and the default, because changing it changes every + * consumer's idea of a tensor's shape. + */ + AS_STORED, + + /** + * Logical `[out, in]`: the convention the engine, HF checkpoints and every kernel assume, and + * the one the block relayout needs. Reverses a 2-D weight's dimensions at the load boundary; + * the bytes are untouched, because they already are `[out, in]` row-major. + */ + OUT_IN, +} 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 05e935135..1e13759dd 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 @@ -6,6 +6,7 @@ import sk.ainet.io.RandomAccessSource import sk.ainet.io.gguf.dequant.DequantOps import sk.ainet.io.model.QuantPolicy import sk.ainet.io.model.StagingPolicy +import sk.ainet.io.model.WeightOrientation import sk.ainet.io.openMappedFile import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor @@ -84,6 +85,18 @@ public class StreamingGgufParametersLoader( * source is not a file, so a browser build behaves exactly as before. */ private val staging: StagingPolicy = StagingPolicy.HEAP, + /** + * Which way round a 2-D weight's shape comes out (#1098, #973 census contradiction #6). + * + * GGUF writes dimensions in `ne` order, so a weight the rest of the engine calls `[out, in]` + * arrives labelled `[in, out]` — while its *bytes* are already `[out, in]` row-major. Nothing + * about the data changes here; only the label. [WeightOrientation.OUT_IN] fixes the label, + * which is what the packed block relayout needs to compute the right permutation. + * + * Defaults to [WeightOrientation.AS_STORED], today's behaviour, because reversing shapes + * changes what every consumer sees. New code should ask for `OUT_IN`. + */ + private val weightOrientation: WeightOrientation = WeightOrientation.AS_STORED, ) : ParametersLoader { init { @@ -94,6 +107,17 @@ public class StreamingGgufParametersLoader( } } + /** + * The shape this loader reports for [tensorInfo], honouring [weightOrientation]: a 2-D weight + * 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 { + val dims = tensorInfo.shape.map { it.toInt() } + val ordered = if (weightOrientation == WeightOrientation.OUT_IN && dims.size == 2) dims.reversed() else dims + return Shape(*ordered.toIntArray()) + } + @Suppress("UNCHECKED_CAST") override suspend fun load( ctx: ExecutionContext, @@ -112,7 +136,7 @@ public class StreamingGgufParametersLoader( var current = 0L for (tensorInfo in tensors) { - val shape = Shape(*tensorInfo.shape.map { it.toInt() }.toIntArray()) + val shape = shapeOf(tensorInfo) // 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). diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/SyntheticGguf.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/SyntheticGguf.kt index edd88165d..0a76e66dd 100644 --- a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/SyntheticGguf.kt +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/SyntheticGguf.kt @@ -22,6 +22,12 @@ object SyntheticGguf { val type: GGMLQuantizationType, val elementCount: Long, val data: ByteArray, + /** + * Dimensions in GGUF `ne` order (fastest-varying first). Defaults to rank 1, which keeps + * element order unambiguous; a 2-D weight is written `[in, out]`, as a real file has it + * (#1098). + */ + val dims: List = listOf(elementCount), ) /** Bytes per block / elements per block for [type], from [GGML_QUANT_SIZES]. */ @@ -165,8 +171,8 @@ object SyntheticGguf { val name = t.name.encodeToByteArray() head.putLong(name.size.toLong()) head.put(name) - head.putInt(1) // rank 1 keeps element order unambiguous - head.putLong(t.elementCount) + head.putInt(t.dims.size) + for (d in t.dims) head.putLong(d) head.putInt(t.type.value) head.putLong(dataOffset) // every payload here is already a multiple of 32 bytes or padded below diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightOrientationTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightOrientationTest.kt new file mode 100644 index 000000000..c38d83346 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightOrientationTest.kt @@ -0,0 +1,79 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.io.model.WeightOrientation +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.Shape +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 + +/** + * #1098 (#973 census contradiction #6): GGUF writes dimensions in `ne` order, so a weight the + * engine calls `[out, in]` arrives labelled `[in, out]` — while its bytes are already `[out, in]` + * row-major. Only the label is wrong, and the label is what the block relayout reads. + * + * The guard that refuses a wrongly-labelled weight lives with the relayout, and is tested in + * `PackedWeightsTest`. + */ +@OptIn(ExperimentalMemoryApi::class) +class WeightOrientationTest { + + /** A Q8_0 weight whose two dimensions differ and where only one of them is block-aligned. */ + private fun file(): File = SyntheticGguf.write( + // ne = [in=128, out=3] → 384 elements, 12 blocks of 32 + SyntheticGguf.tensor("blk.0.attn_q.weight", GGMLQuantizationType.Q8_0, elements = 384), + ) + + private fun load(f: File, orientation: WeightOrientation): Map> { + val ctx = DefaultDataExecutionContext() + val out = LinkedHashMap>() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightOrientation = orientation, + ).load(ctx, FP32::class) { name, t -> out[name] = t } + } + return out + } + + @Test + fun `the default still reports the file's own ne order`() { + val f = SyntheticGguf.write(SyntheticGguf.tensor("w", GGMLQuantizationType.F32, elements = 12)) + try { + // a 1-D tensor has no orientation to get wrong, and nothing changes for it either way + assertEquals(Shape(12), load(f, WeightOrientation.AS_STORED).getValue("w").shape) + assertEquals(Shape(12), load(f, WeightOrientation.OUT_IN).getValue("w").shape) + } finally { + f.delete() + } + } + + @Test + fun `OUT_IN reverses a 2-D weight's label and nothing else`() { + val f = twoDimensionalFile() + try { + val asStored = load(f, WeightOrientation.AS_STORED).getValue("w") + val outIn = load(f, WeightOrientation.OUT_IN).getValue("w") + assertEquals(Shape(64, 4), asStored.shape, "ne order: [in, out]") + assertEquals(Shape(4, 64), outIn.shape, "logical order: [out, in]") + assertContentEquals( + asStored.data.copyToFloatArray(), outIn.data.copyToFloatArray(), + "the bytes are identical — only the label differs", + ) + } finally { + f.delete() + } + } + + /** `ne = [64, 4]`: 256 elements, 8 blocks of 32 — `in = 64`, `out = 4`. */ + private fun twoDimensionalFile(): File { + val t = SyntheticGguf.tensor("w", GGMLQuantizationType.Q8_0, elements = 256) + return SyntheticGguf.write(t.copy(dims = listOf(64L, 4L))) + } +}