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