fix(cpu): stop relayouting a packed weight for kernels this backend does not have

skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api

@@ -273,6 +273,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o
 	protected final fun mapIndex ([ILsk/ainet/lang/tensor/Shape;)[I
 	public fun matmul (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	public fun matmulWeightTransposed (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
+	protected final fun matmulWeightTransposedViaViews (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	public fun maxPool2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;)Lsk/ainet/lang/tensor/Tensor;
 	public fun mean (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;)Lsk/ainet/lang/tensor/Tensor;
 	public fun mulScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor;

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt

@@ -918,14 +918,53 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
     @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 || !isHeapPackedWeight(weight.data)) return matmul(x, transpose(weight))
-        val packed = weight.data as sk.ainet.lang.tensor.storage.PackedBlockStorage
-        val source = packed.packedData
-        val cached = prepackedWeights.firstOrNull { it.first === source }?.second
-        val kernelOrder = cached ?: relayoutPackedWeightForKernels(weight).also { relayouted ->
-            if (prepackedWeights.size >= PREPACK_CACHE_LIMIT) prepackedWeights.removeAt(0)
-            prepackedWeights.add(Pair(source, relayouted as Tensor<*, *>))
-        }
-        return matmul(x, kernelOrder as Tensor<T, V>)
+        // Decode the weight where it lies (#1124). The relayout below produces bytes for kernels
+        // that address `packedData` in feed order themselves; this implementation has no such
+        // kernel, so relayouting for it was pure harm — the result is a tensor whose shape says
+        // [in, out] while its blocks still run along the original input dimension, which no block
+        // order can describe, and decoding it read the wrong blocks and returned plausible garbage.
+        // `KernelDispatch.matmul` already wants the weight output-major, which is exactly the shape
+        // this weight has, so the canonical view goes straight in with nothing rearranged.
+        matmulWeightTransposedViaViews(x, weight)?.let { return it }
+        return matmul(x, transposePackedWeight(weight) ?: return matmulGeneric(x, transpose(weight)))
+    }
+
+    /**
+     * `x · Wᵀ` with [weight] as `[out, in]`, computed by decoding through views (#1124).
+     *
+     * Correct for any packed encoding, because the reference kernel reads through the decoding
+     * `get()`; slower than a packed kernel, which is why [DefaultCpuOpsJvm] overrides this with the
+     * relayout-and-cache path its vectorized kernels can use. `null` when the operands cannot
+     * describe themselves as views.
+     */
+    @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class)
+    @Suppress("UNCHECKED_CAST")
+    protected fun <T : DType, V> matmulWeightTransposedViaViews(x: Tensor<T, V>, weight: Tensor<T, V>): Tensor<T, V>? {
+        if (!sk.ainet.backend.api.kernel.DispatchMode.useRegistry()) return null
+        if (x.dtype != FP32::class) return null
+        val xView = x.data.view ?: return null
+        val wView = weight.data.view ?: return null
+        val (xNorm, leading) = try {
+            sk.ainet.backend.api.kernel.KernelDispatch.normalizeActivation(xView)
+        } catch (_: IllegalArgumentException) {
+            return null
+        }
+        val m = xNorm.shape[0]
+        val k = xNorm.shape[1]
+        val n = wView.shape[0]                       // weight is [out, in]; the dispatcher wants it that way
+        if (k != wView.shape[1]) return null
+        val outArray = FloatArray(m * n)
+        val outView = sk.ainet.lang.memory.TensorView.dense(
+            sk.ainet.lang.memory.Storage.Heap.wrap(outArray), Shape(m, n), FP32,
+        )
+        sk.ainet.backend.api.kernel.KernelDispatch.matmul(xNorm, wView, outView)
+        val outShape = when {
+            x.shape.rank == 1 -> Shape(n)
+            leading.isEmpty() -> Shape(m, n)
+            else -> Shape(*(leading + n))
+        }
+        val outData = dataFactory.fromFloatArray<T, Float>(outShape, x.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData<T, V>
+        return newTensor(outData, x.dtype, x, weight)
     }
 
 

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulEmptyRegistryTest.kt

@@ -0,0 +1,105 @@
+package sk.ainet.exec.tensor.ops
+
+import sk.ainet.backend.api.kernel.KernelDispatch
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.backend.api.kernel.KernelRegistry
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.data.Q4_0BlockTensorData
+import sk.ainet.lang.tensor.data.Q5_0BlockTensorData
+import sk.ainet.lang.tensor.data.Q5_1BlockTensorData
+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.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertTrue
+
+/**
+ * #1124: packed matmul must be right when **nothing** is registered.
+ *
+ * This configuration — the common `DefaultCpuOps` with an empty [KernelRegistry] — was the one
+ * combination no test in the tree covered. The JVM tests always resolve to `DefaultCpuOpsJvm`,
+ * whose override is correct; the native, JS and Wasm tests do use the common implementation, but
+ * their platform factories register `ScalarKernelProvider` first. So every packed-matmul
+ * correctness test, including the ones added for #973, #1096 and #1108, validated a configuration
+ * that was not the broken one.
+ *
+ * The fallback exists precisely for when nothing else is available, which is exactly when it was
+ * wrong — and wrong silently, by 4× the magnitude of the answer.
+ */
+class PackedMatmulEmptyRegistryTest {
+
+    private val ctx = DirectCpuExecutionContext()
+
+    @BeforeTest
+    fun emptyTheRegistry() {
+        // Build the context first: platform factories register ScalarKernelProvider on
+        // construction, and this test is about what happens with nothing registered.
+        ctx.ops
+        KernelRegistry.clearForTesting()
+        KernelDispatch.clearForTesting()
+    }
+
+    @AfterTest
+    fun cleanup() {
+        KernelRegistry.clearForTesting()
+        KernelDispatch.clearForTesting()
+    }
+
+    /** Blocks whose codes differ per block, with a valid fp16 scale of 1.0 (and min 0.0 for Q5_1). */
+    private fun bytes(name: String, blocks: Int, bytesPerBlock: Int): ByteArray {
+        val out = ByteArray(blocks * bytesPerBlock)
+        var seed = 7
+        for (b in 0 until blocks) {
+            val base = b * bytesPerBlock
+            out[base] = 0x00; out[base + 1] = 0x3C
+            if (name == "Q5_1") { out[base + 2] = 0x00; out[base + 3] = 0x00 }
+            for (i in 4 until bytesPerBlock) {
+                seed = seed * 1103515245 + 12345
+                out[base + i] = ((seed ushr 16) % 9 - 4).toByte()
+            }
+        }
+        return out
+    }
+
+    @Test
+    fun `packed matmul agrees with the weight's own decoder when no kernel is registered`() {
+        val cases: List<Triple<String, Int, (Shape, ByteArray) -> TensorData<FP32, Float>>> = listOf(
+            Triple("Q8_0", 34) { s, b -> Q8_0BlockTensorData(s, b) as TensorData<FP32, Float> },
+            Triple("Q4_0", 18) { s, b -> Q4_0BlockTensorData(s, b) as TensorData<FP32, Float> },
+            Triple("Q5_0", 22) { s, b -> Q5_0BlockTensorData(s, b) as TensorData<FP32, Float> },
+            Triple("Q5_1", 24) { s, b -> Q5_1BlockTensorData(s, b) as TensorData<FP32, Float> },
+        )
+        // Three blocks per row, so canonical and kernel-feed order differ (#968), and an output
+        // dimension that is a whole number of blocks so the relayout is row-block-aligned.
+        for ((name, bytesPerBlock, build) in cases) {
+            for ((rows, cols) in listOf(32 to 96, 64 to 96, 32 to 128)) {
+                val w: Tensor<FP32, Float> = ctx.fromData(
+                    build(Shape(rows, cols), bytes(name, rows * (cols / 32), bytesPerBlock)), FP32::class,
+                )
+                val xs = FloatArray(cols) { (it % 13) * 0.0625f }
+                val x = ctx.fromFloatArray<FP32, Float>(Shape(1, cols), FP32::class, xs)
+
+                val decoded = (w.data as PackedBlockStorage).toFloatArray()
+                val expected = FloatArray(rows) { o ->
+                    var acc = 0.0
+                    for (i in 0 until cols) acc += xs[i].toDouble() * decoded[o * cols + i]
+                    acc.toFloat()
+                }
+                val actual = x.matmulWeightTransposed(w).data.copyToFloatArray()
+
+                for (o in 0 until rows) {
+                    assertTrue(
+                        abs(expected[o] - actual[o]) <= 1e-3f * maxOf(1.0f, abs(expected[o])),
+                        "$name [$rows x $cols] output[$o]: expected ${expected[o]}, got ${actual[o]}",
+                    )
+                }
+            }
+        }
+    }
+}

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt

@@ -180,6 +180,39 @@ internal class DefaultCpuOpsJvm(
         return super.divide(a, b)
     }
 
+    /**
+     * `x · Wᵀ` with the weight relayouted once into the order this backend's packed kernels read
+     * (#1096), which is the fast path the common implementation cannot take.
+     *
+     * `DefaultCpuOpsBase` decodes the weight where it lies instead, because it has no kernel that
+     * addresses `packedData` in feed order and relayouting for one that does not exist produced a
+     * tensor nothing could decode correctly (#1124). Here the kernels do exist, so the relayout
+     * pays: once per weight, cached by the identity of its bytes.
+     */
+    @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)
+        val packed = weight.data as sk.ainet.lang.tensor.storage.PackedBlockStorage
+        val source = packed.packedData
+        val cached = prepackedWeightsJvm.firstOrNull { it.first === source }?.second
+        val kernelOrder = cached ?: relayoutPackedWeightForKernels(weight).also { relayouted ->
+            if (prepackedWeightsJvm.size >= PREPACK_CACHE_LIMIT_JVM) prepackedWeightsJvm.removeAt(0)
+            prepackedWeightsJvm.add(Pair(source, relayouted as Tensor<*, *>))
+        }
+        return matmul(x, kernelOrder as Tensor<T, V>)
+    }
+
+    /** Relayouted weights keyed by the identity of the bytes they came from (#1096). */
+    private val prepackedWeightsJvm: MutableList<Pair<ByteArray, Tensor<*, *>>> = mutableListOf()
+    private val PREPACK_CACHE_LIMIT_JVM: Int = 64
+
+    /** The heap packed types whose JVM kernels read input-block-major bytes. */
+    private fun isHeapPackedWeightForJvm(data: sk.ainet.lang.tensor.data.TensorData<*, *>): Boolean =
+        data is sk.ainet.lang.tensor.data.Q4_KTensorData || data is sk.ainet.lang.tensor.data.Q5_KTensorData ||
+            data is sk.ainet.lang.tensor.data.Q6_KTensorData || data is sk.ainet.lang.tensor.data.Q5_1TensorData ||
+            data is sk.ainet.lang.tensor.data.Q5_0TensorData || data is sk.ainet.lang.tensor.data.Q8_0TensorData ||
+            data is sk.ainet.lang.tensor.data.Q4_0TensorData
+
     override fun <T : DType, V> matmul(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
         // `x · Wᵀ` written as two steps (#1108) — before everything, including the strictness check
         // below, which would otherwise report a missing kernel for an operand that has a perfectly