feat(io): weight orientation at the load boundary, and a guard that refuses a transposed label (#973.5)

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.

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

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<IllegalArgumentException> {
+            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<IllegalArgumentException> {
+            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

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,
+}

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 <T : DType, V> 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).

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<Long> = 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

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<String, Tensor<FP32, Float>> {
+        val ctx = DefaultDataExecutionContext()
+        val out = LinkedHashMap<String, Tensor<FP32, Float>>()
+        runBlocking {
+            StreamingGgufParametersLoader(
+                sourceProvider = { JvmRandomAccessSource.open(f) },
+                weightOrientation = orientation,
+            ).load<FP32, Float>(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)))
+    }
+}