From b7a75c9fdb99822c9a41fc67fa9feca0c8cf7c91 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 20:18:26 +0200 Subject: [PATCH] feat(backend): one engine-owned block relayout, and contract fixtures a downstream repo can run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1097 (#973, proposal items 3 and 6). The census in #973 found the same permutation reimplemented in several places and drifting — an inlined Apertus relayout that had diverged from the shared packer it was copied from, three conventions applied to the same GGUF K-quant tensor depending on the converter, and layout knowledge owned by the wrong repository. - `PackedWeights`: `prepackForMatmul` / `toCanonical` for views and `toKernelOrder` / `toCanonicalOrder` for a converter that holds bytes — the only sanctioned implementation of `out[b * rows + o] = in[o * blocksPerRow + b]`. Idempotent by construction: prepacking an already-prepacked weight returns it unchanged, unlike the transpose it replaces, whose double application silently produced a different matrix. `blocksPerRow(encoding, inputDim)` reads the geometry from the encoding's own descriptor rather than a caller's constant. - `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 in the wrong place. Every fixture is three blocks wide, because at one block per row the two orders coincide — the shape that hid #968. - The normative doc gains the downstream section, the "prepack once, at load" rule with the M1-A3 regression that chose the dispatcher's default, and an updated status. Tests: the byte relayout puts each block where the contract says and is its own inverse; every covered format's fixture agrees with its own descriptor; a disagreement names the block that moved; `prepackForMatmul` is idempotent and preserves the decoded matrix; and it refuses what it cannot relayout. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- docs/design/memory/packed-weight-layout.md | 52 ++++++-- .../api/kernel/PackedLayoutFixtures.kt | 108 +++++++++++++++++ .../ainet/backend/api/kernel/PackedWeights.kt | 105 ++++++++++++++++ .../backend/api/kernel/PackedWeightsTest.kt | 113 ++++++++++++++++++ 4 files changed, 371 insertions(+), 7 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedLayoutFixtures.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedWeights.kt create mode 100644 skainet-backends/skainet-backend-api/src/commonTest/kotlin/sk/ainet/backend/api/kernel/PackedWeightsTest.kt diff --git a/docs/design/memory/packed-weight-layout.md b/docs/design/memory/packed-weight-layout.md index 77e723ef2..bb7aa0911 100644 --- a/docs/design/memory/packed-weight-layout.md +++ b/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. diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedLayoutFixtures.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedLayoutFixtures.kt new file mode 100644 index 000000000..200b9eac1 --- /dev/null +++ b/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 = 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 + } +} 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 new file mode 100644 index 000000000..53cefa47e --- /dev/null +++ b/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 + } +} 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 new file mode 100644 index 000000000..904a12ce9 --- /dev/null +++ b/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 { PackedWeights.prepackForMatmul(dense) } + assertFailsWith { PackedWeights.blocksPerRow(TensorEncoding.Dense(4), 128) } + assertFailsWith { + 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 + } +}