feat(io): load a weight in kernel feed order, and have it say so

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;

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 <T : DType, V> rewrapFeedOrderWeight(tensor: Tensor<T, V>): Tensor<T, V>? {
+        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<T, V> = when (tensor.data) {
+            is Q4_KTensorData -> Q4_KBlockTensorData(swapped, bytes) as TensorData<T, V>
+            is Q5_KTensorData -> Q5_KBlockTensorData(swapped, bytes) as TensorData<T, V>
+            is Q6_KTensorData -> Q6_KBlockTensorData(swapped, bytes) as TensorData<T, V>
+            is Q5_1TensorData -> Q5_1BlockTensorData(swapped, bytes) as TensorData<T, V>
+            is Q5_0TensorData -> Q5_0BlockTensorData(swapped, bytes) as TensorData<T, V>
+            is Q8_0TensorData -> Q8_0BlockTensorData(swapped, bytes) as TensorData<T, V>
+            is Q4_0TensorData -> Q4_0BlockTensorData(swapped, bytes) as TensorData<T, V>
+            else -> return null
+        }
+        return newTensor(relabelled, tensor.dtype, tensor)
+    }
+
     @Suppress("UNCHECKED_CAST")
     private fun <T : DType, V> transposePackedWeight(tensor: Tensor<T, V>): Tensor<T, V>? {
         val rank = tensor.shape.rank

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<FP32, Float> =
+        ctx.fromData(
+            Q8_0BlockTensorData(Shape(outDim, inDim), payload, order) as TensorData<FP32, Float>,
+            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<FP32, Float>(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<FP32, Float>(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]}",
+            )
+        }
+    }
+}

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 <T : DType, V> matmulWeightTransposed(x: Tensor<T, V>, weight: Tensor<T, V>): Tensor<T, V> {
         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

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<T, V>, 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<T, V>, 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 {

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<Tensor<FP32, Float>, RecordingTraceSink> {
+        val ctx = DirectCpuExecutionContext()
+        var w: Tensor<FP32, Float>? = null
+        runBlocking {
+            StreamingGgufParametersLoader(
+                sourceProvider = { JvmRandomAccessSource.open(f) },
+                weightForm = form,
+                traceSink = sink,
+            ).load<FP32, Float>(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<FP32, Float>(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<TraceEvent.AdapterInserted>()
+                .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<TraceEvent.AdapterInserted>().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<IllegalArgumentException> {
+            StreamingGgufParametersLoader(
+                sourceProvider = { JvmRandomAccessSource.open(modelFile()) },
+                weightForm = WeightForm(order = WeightByteOrder.KERNEL_FEED),
+            )
+        }
+        assertTrue(failure.message!!.contains("OUT_IN"), failure.message!!)
+    }
+}