From 95d953a035d9b8c066c6db765298245ee59b0cb2 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 25 Aug 2026 18:03:31 +0200 Subject: [PATCH] feat(io): load a weight in kernel feed order, and have it say so MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1120. Follows #1124/#1126. A weight can now be permuted into the order the packed matmul kernels read at load time, and the resulting TensorData declares BlockOrder.INPUT_BLOCK_MAJOR rather than pretending to be canonical. Every reader is then right about the same bytes: the kernels address them in feed order deliberately, while toFloatArray() and view decoding walk the logical grid and fetch each block from where that order put it. The payoff is that #1096's O(bytes) per-weight relayout stops running at all for such a weight. DefaultCpuOpsJvm relabels the shape over the same array instead — asserted structurally with assertSame on the byte array, since a relayout would necessarily produce a different one, and timing proves nothing. Two things that had been conflated are now separate. relayoutPackedWeightForKernels still produces bytes labelled [in, out] for the JVM kernels; that is a private artifact whose shape and blocks disagree, which is exactly why #1126 found it undecodable and why it keeps no declared order. A loader-produced feed-order weight is a different thing: it keeps its [out, in] shape and describes itself. The permutation reuses TensorView.prepack(INPUT_BLOCK_MAJOR), which already owned it and already emits the conversion on the trace — so #1117 reports the prepack without new code, and a test counts exactly one per weight and none for a weight kept as stored. KERNEL_FEED requires WeightShapeOrientation.OUT_IN, and says so: which block is "block b of output row o" has no answer while the tensor is still labelled in the file's ne order, so refusing beats permuting against a meaningless grid. #1115's test asserting the old refusal is rewritten rather than deleted — it now pins the constraint that remains. Gate: scripts/pr-gate.sh — all legs passed. One earlier run failed SlicingTest.testPerformanceAccessPatterns on ChromeHeadless, an unrelated dense-slicing test that passes on its own; not touched here. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-backend-cpu.api | 1 + .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 30 ++++ .../tensor/ops/FeedOrderWeightNoCopyTest.kt | 100 +++++++++++ .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 4 + .../io/gguf/StreamingGgufParametersLoader.kt | 47 ++++- .../sk/ainet/io/gguf/KernelFeedOrderTest.kt | 163 ++++++++++++++++++ .../io/gguf/WeightFormLoaderParityTest.kt | 20 ++- .../lang/tensor/storage/PackedBlockStorage.kt | 15 ++ 8 files changed, 367 insertions(+), 13 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 6757ae6ee..b6594f7a7 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -289,6 +289,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; protected final fun requireSameDType (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; + protected final fun rewrapFeedOrderWeight (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun scaledDotProductAttention (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZ)Lsk/ainet/lang/tensor/Tensor; public fun sigmoid (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index d19d8e384..15e2a531a 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -1002,6 +1002,36 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory "no packed relayout for ${weight.data::class.simpleName}", ) + /** + * A weight already stored in kernel feed order, relabelled as the `[in, out]` tensor the packed + * kernels take — **sharing the same bytes** (#1120). + * + * This is the payoff of letting packed storage declare its order. `transposePackedWeight` is an + * O(bytes) permutation run once per weight and cached; when the loader already produced feed + * order there is nothing to permute, and all that is needed is the other shape label over the + * same array. `null` when [tensor] is not a feed-order packed weight. + */ + @Suppress("UNCHECKED_CAST") + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + protected fun rewrapFeedOrderWeight(tensor: Tensor): Tensor? { + if (tensor.shape.rank != 2) return null + val packed = tensor.data as? sk.ainet.lang.tensor.storage.PackedBlockStorage ?: return null + if (packed.blockOrder != sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR) return null + val swapped = Shape(tensor.shape[1], tensor.shape[0]) + val bytes = packed.packedData + val relabelled: TensorData = when (tensor.data) { + is Q4_KTensorData -> Q4_KBlockTensorData(swapped, bytes) as TensorData + is Q5_KTensorData -> Q5_KBlockTensorData(swapped, bytes) as TensorData + is Q6_KTensorData -> Q6_KBlockTensorData(swapped, bytes) as TensorData + is Q5_1TensorData -> Q5_1BlockTensorData(swapped, bytes) as TensorData + is Q5_0TensorData -> Q5_0BlockTensorData(swapped, bytes) as TensorData + is Q8_0TensorData -> Q8_0BlockTensorData(swapped, bytes) as TensorData + is Q4_0TensorData -> Q4_0BlockTensorData(swapped, bytes) as TensorData + else -> return null + } + return newTensor(relabelled, tensor.dtype, tensor) + } + @Suppress("UNCHECKED_CAST") private fun transposePackedWeight(tensor: Tensor): Tensor? { val rank = tensor.shape.rank diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt new file mode 100644 index 000000000..5fd8712a3 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt @@ -0,0 +1,100 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * #1120: a weight the loader already put in kernel feed order costs nothing to use. + * + * The point of declaring block order is not tidiness — it is that the O(bytes) relayout #1096 runs + * once per weight stops running at all. Asserted structurally rather than by timing: the bytes the + * kernels are handed must be *the same array*, not an equal one. + */ +@OptIn(ExperimentalMemoryApi::class) +class FeedOrderWeightNoCopyTest { + + private val ctx = DirectCpuExecutionContext() + private val outDim = 32 + private val inDim = 96 // three blocks per row: the orders differ (#968) + + private fun bytes(): ByteArray { + val out = ByteArray(outDim * (inDim / 32) * 34) + var seed = 7 + for (b in 0 until outDim * (inDim / 32)) { + val base = b * 34 + out[base] = 0x00; out[base + 1] = 0x3C + for (i in 0 until 32) { + seed = seed * 1103515245 + 12345 + out[base + 2 + i] = ((seed ushr 16) % 9 - 4).toByte() + } + } + return out + } + + @Suppress("UNCHECKED_CAST") + private fun weight(order: BlockOrder, payload: ByteArray): Tensor = + ctx.fromData( + Q8_0BlockTensorData(Shape(outDim, inDim), payload, order) as TensorData, + FP32::class, + ) + + @Test + fun `a feed-order weight reaches the kernels without its bytes being copied`() { + val payload = bytes() + val w = weight(BlockOrder.INPUT_BLOCK_MAJOR, payload) + + val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, FloatArray(inDim) { (it % 13) * 0.0625f }) + val result = x.matmulWeightTransposed(w) + assertEquals(Shape(1, outDim), result.shape) + + // The claim: the same array, still. A relayout would have produced a different one. + assertSame(payload, (w.data as PackedBlockStorage).packedData, "the weight's own bytes must be untouched") + } + + @Test + fun `a feed-order weight and the canonical weight it came from give the same product`() { + val canonical = bytes() + val canonicalWeight = weight(BlockOrder.ROW_MAJOR, canonical) + + // Permute canonically-ordered blocks into feed order by hand: block (o, b) moves from + // o * blocksPerRow + b to b * rows + o. + val blocksPerRow = inDim / 32 + val feed = ByteArray(canonical.size) + for (o in 0 until outDim) { + for (b in 0 until blocksPerRow) { + canonical.copyInto(feed, (b * outDim + o) * 34, (o * blocksPerRow + b) * 34, (o * blocksPerRow + b + 1) * 34) + } + } + val feedWeight = weight(BlockOrder.INPUT_BLOCK_MAJOR, feed) + + assertTrue(!canonical.contentEquals(feed), "the two orders must differ here or the test is vacuous") + assertTrue( + (canonicalWeight.data as PackedBlockStorage).toFloatArray() + .contentEquals((feedWeight.data as PackedBlockStorage).toFloatArray()), + "different bytes, same matrix", + ) + + val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, FloatArray(inDim) { (it % 13) * 0.0625f }) + val fromCanonical = x.matmulWeightTransposed(canonicalWeight).data.copyToFloatArray() + val fromFeed = x.matmulWeightTransposed(feedWeight).data.copyToFloatArray() + for (o in fromCanonical.indices) { + assertTrue( + abs(fromCanonical[o] - fromFeed[o]) <= 1e-3f * maxOf(1.0f, abs(fromCanonical[o])), + "output[$o]: canonical ${fromCanonical[o]} vs feed ${fromFeed[o]}", + ) + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 0f5992404..e62dde190 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -192,6 +192,10 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") override fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { if (weight.shape.rank != 2 || !isHeapPackedWeightForJvm(weight.data)) return super.matmulWeightTransposed(x, weight) + // Already in feed order — the loader produced it that way (#1120). Nothing to permute and + // nothing to cache: the kernels want the other shape label over the same bytes, which costs + // an object rather than a copy of the weight. + rewrapFeedOrderWeight(weight)?.let { return matmul(x, it) } val packed = weight.data as sk.ainet.lang.tensor.storage.PackedBlockStorage val source = packed.packedData val cached = prepackedWeightsJvm.firstOrNull { it.first === source }?.second 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 b84293597..f9f3dabbd 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 @@ -210,12 +210,10 @@ public class StreamingGgufParametersLoader( "ignored; pass only the form." } } - require(form.order == WeightByteOrder.AS_STORED) { - "WeightByteOrder.KERNEL_FEED is not supported by this loader yet (#1120). The bytes are " + - "easy — TensorView.prepack(INPUT_BLOCK_MAJOR) does the permutation — but packed " + - "TensorData addresses packedData as canonical row-major in getBlockScale, getCode " + - "and dequantizeBlock, so feed-order bytes would decode the wrong elements without " + - "failing (#973, #968). Packed storage has to be able to declare its own order first." + require(form.order == WeightByteOrder.AS_STORED || form.shape == WeightShapeOrientation.OUT_IN) { + "WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is defined " + + "relative to a [out, in] weight — which block is 'block b of output row o' has no " + + "answer while the tensor is still labelled in the file's `ne` order." } val requested = form.encoding require(requested !is EncodingRequest.RequantizeTo) { @@ -433,7 +431,42 @@ public class StreamingGgufParametersLoader( "quantizedTensor called for non-quantized type ${tensorInfo.tensorType}" ) } - return ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + val delivered = if (form.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed + return ctx.fromData(delivered as sk.ainet.lang.tensor.data.TensorData, dtype) + } + + /** + * [packed] with its blocks permuted into the order the packed matmul kernels read, and *saying + * so* (#1120). + * + * The permutation is `TensorView.prepack`, which already owns it and emits the conversion on + * the trace (#1117). What #1120 adds is the second half: the result is rebuilt as a + * `TensorData` carrying `BlockOrder.INPUT_BLOCK_MAJOR`, so every reader is right about the same + * bytes — the kernels address them in feed order deliberately, and anything decoding through + * the view or `toFloatArray()` walks the logical grid and fetches each block from where this + * order put it. Before that, feed-order bytes in a type claiming to be canonical decoded to + * plausible garbage (#1124, #973, #968). + */ + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + private fun feedOrdered( + packed: sk.ainet.lang.tensor.storage.PackedBlockStorage, + tensorInfo: StreamingTensorInfo, + ): sk.ainet.lang.tensor.storage.PackedBlockStorage { + val shape = packed.shape + if (shape.rank != 2 || shape[1] % packed.blockSize != 0) return packed + val prepacked = packed.packedView.prepack(sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR, sink = traceSink) + val bytes = (prepacked.storage as? sk.ainet.lang.memory.Storage.Heap)?.bytes ?: return packed + val order = sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR + return when (tensorInfo.tensorType) { + GGMLQuantizationType.Q4_K -> Q4_KBlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q5_K -> Q5_KBlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q6_K -> Q6_KBlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q8_0 -> Q8_0BlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q4_0 -> Q4_0BlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q5_0 -> Q5_0BlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q5_1 -> Q5_1BlockTensorData(shape, bytes, order) + else -> packed + } } private fun bytesToFloatArray(bytes: ByteArray): FloatArray { diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt new file mode 100644 index 000000000..401a47f13 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt @@ -0,0 +1,163 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.exec.kernel.ScalarKernelProvider +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.math.abs +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1120: a weight loaded in kernel feed order says so, decodes to the same matrix, and reaches the + * kernels without the per-weight relayout ever running. + * + * Feed order is only meaningful at three-plus blocks per row — at one block per row the two orders + * coincide and every assertion below would hold vacuously (#968). + */ +@OptIn(ExperimentalMemoryApi::class) +class KernelFeedOrderTest { + + private val outDim = 32 + private val inDim = 96 // three Q8_0 blocks per row + + @BeforeTest + fun registerKernels() { + KernelRegistry.register(ScalarKernelProvider) + } + + private fun modelFile(type: GGMLQuantizationType = GGMLQuantizationType.Q8_0): File = + SyntheticGguf.write( + SyntheticGguf.tensor("blk.0.attn_q.weight", type, elements = outDim * inDim) + .copy(dims = listOf(inDim.toLong(), outDim.toLong())), + ) + + private fun load(f: File, form: WeightForm, sink: RecordingTraceSink = RecordingTraceSink()): + Pair, RecordingTraceSink> { + val ctx = DirectCpuExecutionContext() + var w: Tensor? = null + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightForm = form, + traceSink = sink, + ).load(ctx, FP32::class) { _, t -> w = t } + } + return w!! to sink + } + + private fun asStored() = WeightForm(shape = WeightShapeOrientation.OUT_IN) + private fun feedOrder() = + WeightForm(order = WeightByteOrder.KERNEL_FEED, shape = WeightShapeOrientation.OUT_IN) + + @Test + fun `a feed-order weight declares its order and keeps its shape`() { + val f = modelFile() + try { + val (w, _) = load(f, feedOrder()) + val packed = w.data as PackedBlockStorage + assertEquals(BlockOrder.INPUT_BLOCK_MAJOR, packed.blockOrder, "the bytes must say what order they are in") + assertEquals(Shape(outDim, inDim), w.shape, "feed order is a property of the bytes, not of the shape") + } finally { + f.delete() + } + } + + @Test + fun `feed-order bytes decode to the same matrix they were permuted from`() { + val f = modelFile() + try { + val (canonical, _) = load(f, asStored()) + val (feed, _) = load(f, feedOrder()) + + val canonicalBytes = (canonical.data as PackedBlockStorage).packedData + val feedBytes = (feed.data as PackedBlockStorage).packedData + assertTrue( + !canonicalBytes.contentEquals(feedBytes), + "at three blocks per row the two orders must differ, or this test proves nothing", + ) + + assertContentEquals( + (canonical.data as PackedBlockStorage).toFloatArray(), + (feed.data as PackedBlockStorage).toFloatArray(), + "different bytes, same matrix — that is the whole claim of a declared block order", + ) + } finally { + f.delete() + } + } + + @Test + fun `the product is the same whichever order the weight arrived in`() { + val f = modelFile() + try { + val ctx = DirectCpuExecutionContext() + val xs = FloatArray(inDim) { (it % 13) * 0.0625f } + val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, xs) + + val (canonical, _) = load(f, asStored()) + val (feed, _) = load(f, feedOrder()) + + val fromCanonical = x.matmulWeightTransposed(canonical).data.copyToFloatArray() + val fromFeed = x.matmulWeightTransposed(feed).data.copyToFloatArray() + + for (o in fromCanonical.indices) { + assertTrue( + abs(fromCanonical[o] - fromFeed[o]) <= 1e-3f * maxOf(1.0f, abs(fromCanonical[o])), + "output[$o]: canonical ${fromCanonical[o]} vs feed-order ${fromFeed[o]}", + ) + } + } finally { + f.delete() + } + } + + @Test + fun `the permutation happens once at load and is reported there`() { + val f = modelFile() + try { + val (_, sink) = load(f, feedOrder()) + val relayouts = sink.events() + .filterIsInstance() + .filter { it.kind.startsWith("prepack") } + assertEquals(1, relayouts.size, "one weight, one permutation: ${sink.events().map { it }}") + + val (_, quiet) = load(f, asStored()) + assertTrue( + quiet.events().filterIsInstance().none { it.kind.startsWith("prepack") }, + "a weight kept as stored is not permuted", + ) + } finally { + f.delete() + } + } + + @Test + fun `feed order without OUT_IN is refused because it would be meaningless`() { + val failure = assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(modelFile()) }, + weightForm = WeightForm(order = WeightByteOrder.KERNEL_FEED), + ) + } + assertTrue(failure.message!!.contains("OUT_IN"), failure.message!!) + } +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt index 0a4ac1d08..f79490fbb 100644 --- a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt @@ -140,18 +140,26 @@ class WeightFormLoaderParityTest { } @Test - fun `KERNEL_FEED is refused for the reason it is refused`() { - // Not "unsupported": the bytes are the easy part. The refusal is because packed TensorData - // reads packedData as canonical row-major, so feed-order bytes decode wrong silently (#1120). + fun `KERNEL_FEED is accepted now that packed storage can declare its order`() { + // This slice refused it: packed TensorData read packedData as canonical row-major, so + // feed-order bytes decoded to the wrong elements silently. #1120 gave the bytes a way to + // say what order they are in, so the refusal is gone — replaced by the one constraint that + // remains real, since feed order is defined relative to an [out, in] weight. + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file()) }, + weightForm = WeightForm( + order = WeightByteOrder.KERNEL_FEED, + shape = WeightShapeOrientation.OUT_IN, + ), + ) + val failure = assertFailsWith { StreamingGgufParametersLoader( sourceProvider = { JvmRandomAccessSource.open(file()) }, weightForm = WeightForm(order = WeightByteOrder.KERNEL_FEED), ) } - val message = failure.message!! - assertTrue(message.contains("#1120"), "it names where this is being fixed: $message") - assertTrue(message.contains("canonical row-major"), "and why it cannot be faked: $message") + assertTrue(failure.message!!.contains("OUT_IN"), failure.message!!) } @Test diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt index 6cefa7266..868487084 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt @@ -93,6 +93,21 @@ public interface PackedBlockStorage { public fun toFloatArray(): FloatArray { val result = FloatArray(shape.volume) + if (blockOrder == sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR && shape.rank == 2) { + // Feed-order bytes hold block `(o, b)` at physical index `b * rows + o`, so decoding + // them in physical sequence would emit the matrix transposed-in-blocks. Walk the + // *logical* grid instead and fetch each block from where this order put it — which is + // what makes a feed-order weight decode to the same matrix as the canonical one it was + // permuted from (#1120). + val rows = shape[0] + val blocksPerRow = shape[1] / blockSize + for (o in 0 until rows) { + for (b in 0 until blocksPerRow) { + dequantizeBlock(b * rows + o, result, (o * blocksPerRow + b) * blockSize) + } + } + return result + } var offset = 0 for (i in 0 until blockCount) { val remaining = shape.volume - offset