feat(backend): one engine-owned block relayout, and contract fixtures a downstream repo can run (#973.4)

docs/design/memory/packed-weight-layout.md

@@ -53,6 +53,29 @@ wearing transpose's name and a swapped shape label that lied about the data
 (see #973, "the deeper semantic problem"); replacing that with a weight-transposing matmul
 primitive is [#1096](https://github.com/SKaiNET-developers/SKaiNET/issues/1096).
 
+## For a downstream repository
+
+Two things are published with the engine so a converter never has to reimplement this:
+
+- **`PackedWeights`** — `prepackForMatmul(view)` / `toCanonical(view)` for views, and
+  `toKernelOrder(bytes, rows, blocksPerRow, bytesPerBlock)` / `toCanonicalOrder(...)` for a
+  converter that holds bytes. This is the *only* sanctioned implementation of the permutation. A
+  private copy is what #973 exists to stop: the census found one that had drifted from the shared
+  packer it was copied from, and layout knowledge living in the wrong repository.
+- **`PackedLayoutFixtures`** — canonical and kernel-order fixtures per format, in `main` rather than
+  a test source set, so the artifact a downstream repository already depends on carries them.
+  `disagreement(bytes, encoding, kernelOrder)` returns `null` when the bytes agree, or names the
+  first block that is in the wrong place.
+
+A downstream test asserting `disagreement(myConverterOutput, Q4_K, kernelOrder = true) == null` is
+running against the same bytes the engine's own tests run against. That is what makes a layout
+change fail somewhere rather than ship: previously each repository's suite proved only its own
+convention, and neither crossed the boundary — which is how a byte-layout change shipped as a
+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.
+
 ## Rules
 
 1. **A file's bytes are `ROW_MAJOR`.** Anything loaded from GGUF, produced by a quantizer, or
@@ -66,12 +89,27 @@ primitive is [#1096](https://github.com/SKaiNET-developers/SKaiNET/issues/1096).
    minor/major change, never a patch — this is what let a byte-layout change ship as a green-CI
    hotfix once already.
 
+## Prepack once, at load
+
+The relayout is O(bytes). A weight prepacked at load hits the packed kernel's key directly and the
+dispatcher copies nothing per call; a canonical weight handed straight to the dispatcher gets the
+decoding reference kernel — correct, and slower.
+
+`KernelDispatch.matmul(..., prepackWeights = true)` will relayout for you, and it is **off by
+default** because doing it inside a decode step copies the whole weight per token. That is the same
+per-forward copy #973 objects to in `ops.transpose`, merely moved; wiring it on by default broke
+M1-A3 in exactly that way during #1095, which is how the default was chosen.
+
 ## Status
 
-`Layout.blockOrder`, the `LayoutClass` split, `prepack` and this document land with
-[#1094](https://github.com/SKaiNET-developers/SKaiNET/issues/1094). The remaining work is tracked
-as sub-issues of #973: bridging the packed SPI kernels through the ordered key
-([#1095](https://github.com/SKaiNET-developers/SKaiNET/issues/1095)), the weight-transposing matmul
-primitive (#1096), engine-owned prepacking and cross-repo contract fixtures
-([#1097](https://github.com/SKaiNET-developers/SKaiNET/issues/1097)), and normalizing weight
-orientation at the load boundary ([#1098](https://github.com/SKaiNET-developers/SKaiNET/issues/1098)).
+`Layout.blockOrder`, the `LayoutClass` split, `prepack` and this document landed with
+[#1094](https://github.com/SKaiNET-developers/SKaiNET/issues/1094); the packed SPI kernels reach the
+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).
+
+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.

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedLayoutFixtures.kt

@@ -0,0 +1,108 @@
+package sk.ainet.backend.api.kernel
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.blockSpec
+import sk.ainet.lang.tensor.storage.TensorEncoding
+
+/**
+ * The packed block-layout contract as **fixtures another repository can run** (#973 proposal item
+ * 6; #1097).
+ *
+ * A byte-layout change once shipped as a green-CI hotfix because each repository's tests proved
+ * only its own convention: the engine built canonical fixtures, the downstream converter built
+ * kernel-native ones, and neither suite crossed the boundary. These builders are published *with
+ * the engine*, so a downstream test can assert against the same bytes the engine's own tests
+ * assert against — and a change to what "canonical" means fails that test instead of silently
+ * shipping.
+ *
+ * Deliberately in `main`, not a test source set: the whole point is that the artifact a downstream
+ * repository already depends on carries them.
+ *
+ * Every fixture is **three blocks wide** on purpose. At one block per row the two orders coincide,
+ * which is exactly the shape that hid #968.
+ */
+@ExperimentalMemoryApi
+public object PackedLayoutFixtures {
+
+    /** The formats the contract covers — everything with a block geometry and a matmul kernel. */
+    public val encodings: List<TensorEncoding> = 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,
+    )
+
+    /** Default fixture geometry: four output rows × three blocks per row. */
+    public const val ROWS: Int = 4
+    public const val BLOCKS_PER_ROW: Int = 3
+
+    /**
+     * Canonical (`ROW_MAJOR`) bytes for [encoding]: block `(o, b)` at flat index
+     * `o * blocksPerRow + b`, every byte of it set to a value that identifies the block.
+     *
+     * The content is deliberately trivial and identifying rather than realistic: this fixture
+     * exists to pin *where blocks are*, not what they decode to. Decode fidelity is the golden
+     * gate's job.
+     */
+    public fun canonical(
+        encoding: TensorEncoding,
+        rows: Int = ROWS,
+        blocksPerRow: Int = BLOCKS_PER_ROW,
+    ): ByteArray {
+        val bytesPerBlock = bytesPerBlockOf(encoding)
+        val out = ByteArray(rows * blocksPerRow * bytesPerBlock)
+        for (o in 0 until rows) {
+            for (b in 0 until blocksPerRow) {
+                val base = (o * blocksPerRow + b) * bytesPerBlock
+                for (i in 0 until bytesPerBlock) out[base + i] = blockTag(o, b)
+            }
+        }
+        return out
+    }
+
+    /** The same weight in the order the kernels read (`INPUT_BLOCK_MAJOR`): block `(o, b)` at `b * rows + o`. */
+    public fun kernelOrder(
+        encoding: TensorEncoding,
+        rows: Int = ROWS,
+        blocksPerRow: Int = BLOCKS_PER_ROW,
+    ): ByteArray = PackedWeights.toKernelOrder(canonical(encoding, rows, blocksPerRow), rows, blocksPerRow, bytesPerBlockOf(encoding))
+
+    /** The byte every byte of block `(o, b)` carries — an identity a test can assert on. */
+    public fun blockTag(row: Int, block: Int): Byte = (row * 16 + block).toByte()
+
+    /** Bytes per block of [encoding], from its own descriptor. */
+    public fun bytesPerBlockOf(encoding: TensorEncoding): Int {
+        val spec = encoding.blockSpec ?: throw IllegalArgumentException("$encoding is not block-structured")
+        require(!spec.isPerTensor) { "${encoding.name} has no fixed block size; it is a per-tensor encoding" }
+        return spec.bytesPerBlock
+    }
+
+    /**
+     * Check that [bytes] holds the fixture's blocks in the given order — the assertion a downstream
+     * test runs against its own converter's output.
+     *
+     * @return `null` when it agrees, or a description of the first block that is in the wrong place
+     */
+    public fun disagreement(
+        bytes: ByteArray,
+        encoding: TensorEncoding,
+        kernelOrder: Boolean,
+        rows: Int = ROWS,
+        blocksPerRow: Int = BLOCKS_PER_ROW,
+    ): String? {
+        val bytesPerBlock = bytesPerBlockOf(encoding)
+        val required = rows * blocksPerRow * bytesPerBlock
+        if (bytes.size < required) return "expected at least $required bytes, got ${bytes.size}"
+        for (o in 0 until rows) {
+            for (b in 0 until blocksPerRow) {
+                val index = if (kernelOrder) b * rows + o else o * blocksPerRow + b
+                val actual = bytes[index * bytesPerBlock]
+                val expected = blockTag(o, b)
+                if (actual != expected) {
+                    return "block ($o, $b) should be at flat index $index " +
+                        "(${if (kernelOrder) "input-block-major" else "canonical"}), found tag $actual instead of $expected"
+                }
+            }
+        }
+        return null
+    }
+}

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedWeights.kt

@@ -0,0 +1,105 @@
+package sk.ainet.backend.api.kernel
+
+import sk.ainet.lang.memory.BlockOrder
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Scope
+import sk.ainet.lang.memory.blockSpec
+import sk.ainet.lang.memory.TensorView
+import sk.ainet.lang.memory.trace.NoopTraceSink
+import sk.ainet.lang.memory.trace.TraceSink
+import sk.ainet.lang.tensor.storage.TensorEncoding
+
+/**
+ * The engine's own answer to "how do I get a packed weight into the shape a kernel reads"
+ * (#973 proposal item 3; #1097).
+ *
+ * Before this, every converter kept a private copy of the row-major → input-block-major relayout,
+ * and they diverged: the census in #973 found one inlined in a downstream Apertus path that had
+ * drifted from the shared packer it was copied from, three different conventions applied to the
+ * same GGUF K-quant tensor depending on which converter loaded it, and layout knowledge living in
+ * the wrong repository entirely. One function, owned here, is the fix — and the fixtures in
+ * [PackedLayoutFixtures] are how a downstream repository proves it agrees.
+ *
+ * The relayout is O(bytes). Call it **once, at load**; a weight prepacked at load hits the packed
+ * kernel's key directly and the dispatcher copies nothing per call (#1095).
+ */
+@ExperimentalMemoryApi
+public object PackedWeights {
+
+    /**
+     * [weight] in the order the packed matmul kernels read — [BlockOrder.INPUT_BLOCK_MAJOR].
+     *
+     * Returns the view unchanged when it is already in that order, so calling this twice is
+     * harmless: unlike the packed "transpose" it replaces, it is idempotent by construction.
+     *
+     * @throws IllegalArgumentException if [weight] is not a 2-D block-packed view
+     */
+    public fun prepackForMatmul(
+        weight: TensorView,
+        scope: Scope = Scope.Ambient,
+        sink: TraceSink = NoopTraceSink,
+    ): 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}" }
+        return weight.prepack(BlockOrder.INPUT_BLOCK_MAJOR, scope, sink)
+    }
+
+    /** [weight] back in the canonical order a file holds — the inverse of [prepackForMatmul]. */
+    public fun toCanonical(
+        weight: TensorView,
+        scope: Scope = Scope.Ambient,
+        sink: TraceSink = NoopTraceSink,
+    ): TensorView {
+        require(weight.layout.blocked) { "toCanonical takes a block-packed weight, got ${weight.format}" }
+        return weight.prepack(BlockOrder.ROW_MAJOR, scope, sink)
+    }
+
+    /**
+     * The relayout at the byte level, for a converter that holds bytes rather than views:
+     * `[out, in]` blocks in canonical order → kernel order.
+     *
+     * `out[b * rows + o] = in[o * blocksPerRow + b]`, block by block. This is the *only* sanctioned
+     * implementation of that permutation; a private copy is what #973 exists to stop.
+     */
+    public fun toKernelOrder(canonical: ByteArray, rows: Int, blocksPerRow: Int, bytesPerBlock: Int): ByteArray =
+        permute(canonical, rows, blocksPerRow, bytesPerBlock, toKernelOrder = true)
+
+    /** The inverse of [toKernelOrder]: kernel-order bytes back to canonical. */
+    public fun toCanonicalOrder(kernelOrder: ByteArray, rows: Int, blocksPerRow: Int, bytesPerBlock: Int): ByteArray =
+        permute(kernelOrder, rows, blocksPerRow, bytesPerBlock, toKernelOrder = false)
+
+    private fun permute(
+        source: ByteArray,
+        rows: Int,
+        blocksPerRow: Int,
+        bytesPerBlock: Int,
+        toKernelOrder: Boolean,
+    ): ByteArray {
+        require(rows > 0 && blocksPerRow > 0 && bytesPerBlock > 0) { "geometry must be positive" }
+        val required = rows * blocksPerRow * bytesPerBlock
+        require(source.size >= required) {
+            "need $required bytes for a $rows × $blocksPerRow block grid of ${bytesPerBlock}-byte blocks, got ${source.size}"
+        }
+        val out = ByteArray(required)
+        for (o in 0 until rows) {
+            for (b in 0 until blocksPerRow) {
+                val canonical = (o * blocksPerRow + b) * bytesPerBlock
+                val kernel = (b * rows + o) * bytesPerBlock
+                val from = if (toKernelOrder) canonical else kernel
+                val to = if (toKernelOrder) kernel else canonical
+                source.copyInto(out, to, from, from + bytesPerBlock)
+            }
+        }
+        return out
+    }
+
+    /** Block geometry of [encoding] — what a converter needs to call [toKernelOrder]. */
+    public fun blocksPerRow(encoding: TensorEncoding, inputDim: Int): Int {
+        val spec = encoding.blockSpec
+            ?: throw IllegalArgumentException("$encoding is not block-structured")
+        require(inputDim % spec.blockSize == 0) {
+            "${encoding.name} tiles the input dimension in blocks of ${spec.blockSize}; $inputDim is not a multiple"
+        }
+        return inputDim / spec.blockSize
+    }
+}

skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/PackedWeightsTest.kt

@@ -0,0 +1,113 @@
+package sk.ainet.backend.api.kernel
+
+import sk.ainet.lang.memory.BlockOrder
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.PackedBlockDecoder
+import sk.ainet.lang.memory.Storage
+import sk.ainet.lang.memory.TensorView
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.data.Q8_0BlockTensorData
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertNotEquals
+import kotlin.test.assertNull
+import kotlin.test.assertTrue
+
+/**
+ * #1097 (#973.3): one implementation of the block relayout, owned by the engine, and fixtures a
+ * downstream repository can assert against.
+ *
+ * The census in #973 found the same permutation reimplemented in several places, drifting; the
+ * point of these tests is that there is now exactly one, that it is its own inverse, and that a
+ * change to what "canonical" means fails a test rather than shipping.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class PackedWeightsTest {
+
+    private val rows = 4
+    private val blocksPerRow = 3
+    private val bytesPerBlock = 34          // Q8_0
+
+    @Test
+    fun theByteLevelRelayoutMovesEachBlockWhereTheContractSays() {
+        val canonical = PackedLayoutFixtures.canonical(TensorEncoding.Q8_0)
+        val kernelOrder = PackedWeights.toKernelOrder(canonical, rows, blocksPerRow, bytesPerBlock)
+        assertNull(PackedLayoutFixtures.disagreement(canonical, TensorEncoding.Q8_0, kernelOrder = false))
+        assertNull(PackedLayoutFixtures.disagreement(kernelOrder, TensorEncoding.Q8_0, kernelOrder = true))
+        assertNotEquals(
+            canonical.toList(), kernelOrder.toList(),
+            "with three blocks per row the two orders must differ — at one block they coincide, which is what hid #968",
+        )
+    }
+
+    @Test
+    fun theRelayoutIsItsOwnInverse() {
+        val canonical = PackedLayoutFixtures.canonical(TensorEncoding.Q4_K)
+        val bytes = PackedLayoutFixtures.bytesPerBlockOf(TensorEncoding.Q4_K)
+        val roundTrip = PackedWeights.toCanonicalOrder(
+            PackedWeights.toKernelOrder(canonical, rows, blocksPerRow, bytes), rows, blocksPerRow, bytes,
+        )
+        assertContentEquals(canonical, roundTrip, "the permutation must be invertible — unlike the transpose it replaces")
+    }
+
+    @Test
+    fun everyCoveredFormatHasAFixtureThatAgreesWithItsOwnDescriptor() {
+        for (encoding in PackedLayoutFixtures.encodings) {
+            val bytes = PackedLayoutFixtures.bytesPerBlockOf(encoding)
+            assertEquals(
+                bytes.toLong() * blocksPerRow, encoding.physicalBytes(blocksPerRow.toLong() * blockSizeOf(encoding)),
+                "${encoding.name}: the fixture's block size must be the encoding's own",
+            )
+            assertNull(PackedLayoutFixtures.disagreement(PackedLayoutFixtures.canonical(encoding), encoding, kernelOrder = false), encoding.name)
+            assertNull(PackedLayoutFixtures.disagreement(PackedLayoutFixtures.kernelOrder(encoding), encoding, kernelOrder = true), encoding.name)
+        }
+    }
+
+    @Test
+    fun aDisagreementIsReportedWithTheBlockThatMoved() {
+        // what a downstream test sees when its converter emits the other order
+        val wrong = PackedLayoutFixtures.kernelOrder(TensorEncoding.Q8_0)
+        val message = PackedLayoutFixtures.disagreement(wrong, TensorEncoding.Q8_0, kernelOrder = false)
+        assertTrue(message != null && message.contains("should be at flat index"), message ?: "no disagreement reported")
+        assertTrue(message!!.contains("canonical"), message)
+    }
+
+    @Test
+    fun prepackForMatmulIsIdempotent() {
+        val bytes = PackedLayoutFixtures.canonical(TensorEncoding.Q8_0)
+        val shape = Shape(rows, blocksPerRow * 32)
+        val view = TensorView.packed(
+            Storage.Heap.wrap(bytes), shape, TensorEncoding.Q8_0,
+            PackedBlockDecoder(Q8_0BlockTensorData(shape, bytes)),
+        )
+        val once = PackedWeights.prepackForMatmul(view)
+        val twice = PackedWeights.prepackForMatmul(once)
+        assertEquals(BlockOrder.INPUT_BLOCK_MAJOR, once.layout.blockOrder)
+        assertEquals(once.storage.id, twice.storage.id, "calling it twice must not copy again — it is not a transpose")
+        assertContentEquals(view.toFloatArray(), once.toFloatArray(), "and the matrix is unchanged")
+
+        val back = PackedWeights.toCanonical(once)
+        assertEquals(BlockOrder.ROW_MAJOR, back.layout.blockOrder)
+        assertContentEquals(view.toFloatArray(), back.toFloatArray())
+    }
+
+    @Test
+    fun itRefusesWhatItCannotRelayout() {
+        val dense = TensorView.dense(Storage.Heap.floats(16), Shape(4, 4))
+        assertFailsWith<IllegalArgumentException> { PackedWeights.prepackForMatmul(dense) }
+        assertFailsWith<IllegalArgumentException> { PackedWeights.blocksPerRow(TensorEncoding.Dense(4), 128) }
+        assertFailsWith<IllegalArgumentException> {
+            PackedWeights.blocksPerRow(TensorEncoding.Q8_0, 100)   // not a multiple of the block size
+        }
+        assertEquals(4, PackedWeights.blocksPerRow(TensorEncoding.Q8_0, 128))
+        assertEquals(2, PackedWeights.blocksPerRow(TensorEncoding.Q4_K, 512))
+    }
+
+    private fun blockSizeOf(encoding: TensorEncoding): Int = when (encoding) {
+        TensorEncoding.Q4_0, TensorEncoding.Q5_0, TensorEncoding.Q5_1, TensorEncoding.Q8_0 -> 32
+        else -> 256
+    }
+}