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 cedfccd97..0d47b2378 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 @@ -272,6 +272,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun lt (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; 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; 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; @@ -283,6 +284,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun pow (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun powScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public fun rdivScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; 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; 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 cc5e11983..e8623dbdf 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 @@ -872,32 +872,82 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory return newTensor(outData, a.dtype, a, b) } - @TensorOp() - override fun transpose(tensor: Tensor): Tensor { + /** + * Weights already relayouted into kernel feed order, keyed by the identity of the packed bytes + * they came from (#973/#1096). + * + * The relayout is O(bytes). Doing it inside [matmulWeightTransposed] once per weight instead of + * once per call is what removes the per-forward copy `Linear.onForward` used to pay: a model's + * weights are stable, so the first forward pass converts and every later one reuses. Bounded, + * because a cache that grows without limit on a 2 GB device is its own bug; a model with more + * than [PREPACK_CACHE_LIMIT] distinct packed weights simply converts the overflow each time, + * which is exactly the old behaviour. + */ + private val prepackedWeights: MutableList>> = mutableListOf() + + /** How many relayouted weights to keep; beyond this the oldest is dropped and reconverted on demand. */ + private val PREPACK_CACHE_LIMIT: Int = 64 + + /** + * `x · Wᵀ` with the weight as `[out, in]` — the primitive, and the way out of the per-forward + * copy (#973 "the deeper semantic problem", #1096). + * + * For a block-quantized weight this relayouts **once** and reuses the result; for anything else + * it is the ordinary `matmul(x, transpose(w))`, which for dense data is a free shape swap. + */ + @Suppress("UNCHECKED_CAST") + override fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { + 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) + } + + + /** + * The block-grid permutation that used to live in `transpose` (#973/#1096). + * + * The packed matmul kernels read `packedData` as **input-block-major** — + * `(blockIdx * outputDim + o)`, every output row's block for one input block contiguous — + * whatever shape the tensor declares. Canonical packed storage, as loaded from a GGUF or built + * by any row-major producer, is the other order. The two coincide only at one block per row, so + * for a real weight a bare shape relabel hands the kernel bytes in the wrong physical order and + * it reads garbage without failing (#968, and downstream SKaiNET-transformers#307). + * + * So this is a real O(bytes) permutation, and [matmulWeightTransposed] runs it **once** per + * weight rather than once per call — which is the difference #1096 exists to make. + * + * @return the relayouted weight, or `null` for a data type with no packed relayout + */ + /** The heap packed data types whose kernels read input-block-major bytes. */ + private fun isHeapPackedWeight(data: sk.ainet.lang.tensor.data.TensorData<*, *>): Boolean = + data is Q4_KTensorData || data is Q5_KTensorData || data is Q6_KTensorData || + data is Q5_1TensorData || data is Q5_0TensorData || data is Q8_0TensorData || data is Q4_0TensorData + + /** + * The block relayout by its own name (#973/#1096) — what `transpose` used to do to a packed + * weight, for the callers that genuinely want the permuted bytes rather than a product. + * + * Prefer [matmulWeightTransposed], which does this once per weight instead of once per call. + * + * @throws UnsupportedOperationException for a data type with no packed relayout + */ + override fun relayoutPackedWeightForKernels(weight: Tensor): Tensor = + transposePackedWeight(weight) + ?: throw UnsupportedOperationException( + "no packed relayout for ${weight.data::class.simpleName}", + ) + + @Suppress("UNCHECKED_CAST") + private fun transposePackedWeight(tensor: Tensor): Tensor? { val rank = tensor.shape.rank - require(rank >= 2) { "Transpose requires at least 2 dimensions" } val rows = tensor.shape[rank - 2] val cols = tensor.shape[rank - 1] - - // Transpose for heap-packed quant weights (Q4_K/Q5_K/Q6_K/Q5_1/Q5_0/Q8_0/Q4_0): - // the packed-quant matmul kernels (chooseQuantizedMatmulHeap, native/Panama/ - // scalar alike) always read `packedData` as *input-block-major* — - // `(blockIdx * outputDim + o)`, all `outputDim` rows' blocks for a fixed - // input block contiguous (see e.g. q5_0_matmul.c) — regardless of the - // tensor's declared shape. Canonical packed storage (as loaded verbatim from - // GGUF, or as built by any other row-major producer) is instead *row-major*: - // for output row `o`, its `blocksPerInputDim` blocks are contiguous, then row - // `o + 1`. Those two orderings coincide only when there's a single block per - // row (`blocksPerInputDim == 1`); for any weight wider than one block — i.e. - // virtually every real model, since inputDim is almost always >> the 32/256- - // element block size — a bare shape relabel hands the kernel bytes in the - // WRONG physical order and it silently reads garbage (SKaiNET#968, - // surfaced downstream via SKaiNET-transformers#307's all-zero Q5_0/Q5_1 - // matmul). So this performs the actual O(bytes) block-grid permutation - // (`transposePackedBlocks`) instead of a free shape swap. Still avoids the - // FP32 dequant round-trip `ops.matmul(x, ops.transpose(W))` was written to - // dodge — just not for free anymore. - if (rank == 2) { @Suppress("UNCHECKED_CAST") when (val d = tensor.data) { is Q4_KTensorData -> { @@ -941,19 +991,48 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q4_0TensorData.BYTES_PER_BLOCK) return newTensor(Q4_0BlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) } - // Narrow floats (FP16/BF16) relaid input-major at load: the transpose is the - // same buffer read with the other shape's strides, so hand back an ordinary - // dense narrow tensor over it. This is what lets a KEEP_NATIVE weight survive - // `Linear.onForward`'s `weight.t()` and reach the narrow matmul kernel — see - // issue #888. Note the asymmetry with the block-quant arms above: only the - // *input-major* type is safe to reinterpret. A row-major narrow buffer falls - // through to the generic path on purpose, because swapping its shape would - // silently yield a different matrix rather than the transpose. - // Lives only here: DefaultCpuOpsJvm.transpose intercepts nothing that would - // shadow this case, so the JVM falls through to this arm too. - is NarrowFloatInputMajorTensorData -> return newTensor(d.transposedView() as TensorData, tensor.dtype, tensor) else -> {} } + + return null + } + + @TensorOp() + override fun transpose(tensor: Tensor): Tensor { + val rank = tensor.shape.rank + require(rank >= 2) { "Transpose requires at least 2 dimensions" } + val rows = tensor.shape[rank - 2] + val cols = tensor.shape[rank - 1] + + // Only the *heap* packed types, whose kernels read input-block-major bytes. The + // MemorySegment tier reads canonical bytes, so for those a shape swap is genuinely correct + // and stays where it is — census contradiction #3, now stated instead of implied. + val heapPacked = rank == 2 && isHeapPackedWeight(tensor.data) + if (heapPacked) { + val packedData = tensor.data as sk.ainet.lang.tensor.storage.PackedBlockStorage + // Transposing block-quantized data is not a representable operation (#973): blocks + // quantize runs along the input dimension, so a real transpose needs requantization. + // What used to happen here was a layout conversion wearing transpose's name — an + // O(bytes) copy per call, not an involution, and a lie about what the result means. + throw UnsupportedOperationException( + "transpose() is not defined for a ${packedData.encoding.name} weight: blocks quantize runs " + + "along the input dimension, so transposing them needs requantization, and what this used to " + + "do was a per-call layout copy that is not its own inverse (#973). Use " + + "ops.matmulWeightTransposed(x, weight) with the weight as [out, in], or relayout explicitly " + + "with PackedWeights.prepackForMatmul. See docs/design/memory/packed-weight-layout.md.", + ) + } + + // Narrow floats (FP16/BF16) relaid input-major at load are a *view* rewrap, not a packed + // relayout: the transpose is the same buffer read with the other shape's strides, which is + // what lets a KEEP_NATIVE weight survive Linear's weight handling and reach the narrow + // matmul kernel (#888). Only the block-quantized types are refused above (#973). + if (rank == 2) { + val narrow = tensor.data as? NarrowFloatInputMajorTensorData + if (narrow != null) { + @Suppress("UNCHECKED_CAST") + return newTensor(narrow.transposedView() as TensorData, tensor.dtype, tensor) + } } // Fast path: 2D float tensor — direct buffer swap diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/golden/PackedMatmulDispatchParityTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/golden/PackedMatmulDispatchParityTest.kt index 2786a59f2..9bf26db91 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/golden/PackedMatmulDispatchParityTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/golden/PackedMatmulDispatchParityTest.kt @@ -18,7 +18,7 @@ import kotlin.test.assertTrue /** * SKEEP-003 golden gate, dispatch half (runs on every target): for all seven GGML packed - * encodings, `ops.matmul(x, ops.transpose(w))` on a canonical row-major packed weight must agree + * encodings, `ops.matmulWeightTransposed(x, w)` on a canonical row-major packed weight must agree * with an FP32 reference computed from the decoded weight. Tolerance-based because the JVM may * pick SIMD / native / Q8-activation kernel tiers (#944), which are not bit-identical to the * scalar reference; the bit-identical guarantees live in the goldenTest source set. @@ -74,7 +74,7 @@ class PackedMatmulDispatchParityTest { val xf = FloatArray(batch * inDim) { rng.nextFloat() * 2f - 1f } val x = ctx.fromFloatArray(Shape(batch, inDim), FP32::class, xf) - val actual = ctx.ops.matmul(x, ctx.ops.transpose(w)).data.copyToFloatArray() + val actual = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() // FP32 reference from the decoded weight val wf = FloatArray(outDim * inDim); val tmp = FloatArray(f.blockSize) diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/MatmulWeightTransposedTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/MatmulWeightTransposedTest.kt new file mode 100644 index 000000000..12a0b7e89 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/MatmulWeightTransposedTest.kt @@ -0,0 +1,95 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +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.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1096 (#973): `x · Wᵀ` as a primitive, and `transpose` on a packed weight as an error. + * + * The old shape — `matmul(x, transpose(w))` — copied the whole weight on every call, and + * `transpose(transpose(w))` was not `w`. Both are gone: the product is asked for directly, the + * relayout happens once per weight, and `transpose` refuses rather than lying. + */ +class MatmulWeightTransposedTest { + + private val ctx = DirectCpuExecutionContext() + private val outDim = 4 + private val inDim = 96 // three Q8_0 blocks per row: the case where block order matters + + @Suppress("UNCHECKED_CAST") + private fun weight(): Tensor { + val blocks = outDim * (inDim / 32) + val bytes = ByteArray(blocks * 34) + var seed = 11 + for (b in 0 until blocks) { + val base = b * 34 + bytes[base] = 0x00; bytes[base + 1] = 0x3C // fp16 scale 1.0 + for (i in 0 until 32) { + seed = seed * 1103515245 + 12345 + bytes[base + 2 + i] = ((seed ushr 16) % 17 - 8).toByte() + } + } + val data = Q8_0BlockTensorData(Shape(outDim, inDim), bytes) + return ctx.fromData(data as TensorData, FP32::class) + } + + private fun activation(): Tensor = + ctx.fromFloatArray(Shape(1, inDim), FP32::class, FloatArray(inDim) { (it % 7) * 0.125f }) + + @Test + fun `the primitive agrees with the relayout then matmul it replaces`() { + val w = weight() + val x = activation() + val viaPrimitive = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() + val viaRelayout = ctx.ops.matmul(x, ctx.ops.relayoutPackedWeightForKernels(w)).data.copyToFloatArray() + assertContentEquals(viaRelayout, viaPrimitive, "the primitive must compute exactly what the old path did") + } + + @Test + fun `a weight is relayouted once however many times it is used`() { + val w = weight() + val x = activation() + val first = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() + repeat(5) { + assertContentEquals(first, ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray()) + } + // and the answer keeps matching the explicit relayout, so the cache is not stale + assertContentEquals( + ctx.ops.matmul(x, ctx.ops.relayoutPackedWeightForKernels(w)).data.copyToFloatArray(), + first, + ) + } + + @Test + fun `transpose refuses a packed weight and says what to use instead`() { + val failure = assertFailsWith { ctx.ops.transpose(weight()) } + val message = failure.message!! + assertTrue(message.contains("not defined for a Q8_0 weight"), message) + assertTrue(message.contains("matmulWeightTransposed"), "it names the primitive: $message") + assertTrue(message.contains("prepackForMatmul"), "and the explicit relayout: $message") + assertTrue(message.contains("requantization"), "and why: $message") + } + + @Test + fun `a dense weight transposes as it always did`() { + val dense: Tensor = + ctx.fromFloatArray(Shape(outDim, inDim), FP32::class, FloatArray(outDim * inDim) { it * 0.01f }) + val t = ctx.ops.transpose(dense) + assertTrue(t.shape == Shape(inDim, outDim)) + val x = activation() + val viaPrimitive = ctx.ops.matmulWeightTransposed(x, dense).data.copyToFloatArray() + val viaTranspose = ctx.ops.matmul(x, t).data.copyToFloatArray() + for (i in viaPrimitive.indices) { + assertTrue(abs(viaPrimitive[i] - viaTranspose[i]) < 1e-4f, "[$i]: ${viaPrimitive[i]} vs ${viaTranspose[i]}") + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulDispatchTest.kt index 5a1923e45..97ce53739 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulDispatchTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulDispatchTest.kt @@ -18,7 +18,7 @@ import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.FP32 /** - * End-to-end proof that packed-quant weights flow through `ctx.ops.matmul(x, ops.transpose(W))` + * End-to-end proof that packed-quant weights flow through `ctx.ops.matmulWeightTransposed(x, W)` * on EVERY platform — exercising the lazy-transpose shape-swap + `chooseQuantizedMatmulHeap` in * DefaultCpuOpsBase, resolving the registered kernel (scalar on Native/JS/WASM, Panama/FFM on JVM). * Runs on jvmTest AND linuxX64Test; a green linuxX64 run is the headline "Native packed matmul works". @@ -45,7 +45,7 @@ class PackedMatmulDispatchTest { * .transpose` is responsible for the canonical → kernel-native block-grid * permutation (`DefaultCpuOpsBase.transposePackedBlocks`); this generator * must hand it genuinely canonical input or it isn't testing the real - * `ops.matmul(x, ops.transpose(W))` path — see SKaiNET-transformers#307. + * `ops.matmulWeightTransposed(x, W)` path — see SKaiNET-transformers#307. */ private fun q5_1(inDim: Int, outDim: Int, rng: Random): Pair { val blocks = inDim / 32; val bytes = ByteArray(outDim * blocks * 24); val wf = FloatArray(outDim * inDim) @@ -135,7 +135,7 @@ class PackedMatmulDispatchTest { ) val xf = FloatArray(inDim) { rng.nextFloat() - 0.5f } val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, xf) - val out = ctx.ops.matmul(x, ctx.ops.transpose(w)).data.copyToFloatArray() + val out = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() val expected = FloatArray(outDim) { o -> var s = 0f; for (i in 0 until inDim) s += xf[i] * wf[o * inDim + i]; s } var maxErr = 0f; var maxAbs = 1f for (o in 0 until outDim) { maxErr = maxOf(maxErr, abs(expected[o] - out[o])); maxAbs = maxOf(maxAbs, abs(expected[o])) } @@ -173,7 +173,7 @@ class PackedMatmulDispatchTest { val bytes = ByteArray(outDim * (inDim / blockElems) * bpb) val w = ctx.fromData(build(Shape(outDim, inDim), bytes), FP32::class) // The bug threw here for unhandled packed types. - val t = ctx.ops.transpose(w) + val t = ctx.ops.relayoutPackedWeightForKernels(w) assertEquals(Shape(inDim, outDim), t.shape, "$name: transpose did not flip shape") assertTrue( t.data::class.simpleName?.contains("Block") == true, diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt index c553dc029..e3c09c466 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/RegistryMatmulDispatchTest.kt @@ -58,7 +58,7 @@ class RegistryMatmulDispatchTest { // #993: the first post-prefill decode step passes a rank-1 [hidden] activation val (w, wf) = packedWeight(4) val x = ctx.fromFloatArray(Shape(32), FP32::class, FloatArray(32) { (it % 7) * 0.5f }) - val out = ctx.ops.matmul(x, ctx.ops.transpose(w)) // [32] x [32, 4] -> [4] + val out = ctx.ops.matmulWeightTransposed(x, w) // [32] x [32, 4] -> [4] val got = out.data.copyToFloatArray() assertTrue(out.shape.rank == 1 && out.shape[0] == 4, "expected [4], was ${out.shape}") for (j in 0 until 4) { @@ -73,7 +73,7 @@ class RegistryMatmulDispatchTest { fun registryAndLegacyPathsAgreeOnAPackedWeight() { val (w, _) = packedWeight(3) val x = ctx.fromFloatArray(Shape(2, 32), FP32::class, FloatArray(64) { (it % 5) * 0.25f }) - val wt = ctx.ops.transpose(w) + val wt = ctx.ops.relayoutPackedWeightForKernels(w) DispatchMode.overrideEnabled = true val viaRegistry = ctx.ops.matmul(x, wt).data.copyToFloatArray() @@ -90,7 +90,7 @@ class RegistryMatmulDispatchTest { fun batchedActivationsFlattenAndReshape() { val (w, wf) = packedWeight(2) val x = ctx.fromFloatArray(Shape(2, 3, 32), FP32::class, FloatArray(192) { (it % 4).toFloat() }) - val out = ctx.ops.matmul(x, ctx.ops.transpose(w)) + val out = ctx.ops.matmulWeightTransposed(x, w) assertTrue(out.shape.dimensions.toList() == listOf(2, 3, 2), "expected [2, 3, 2], was ${out.shape}") val got = out.data.copyToFloatArray() var expect = 0f @@ -107,7 +107,7 @@ class RegistryMatmulDispatchTest { @Suppress("UNCHECKED_CAST") val w = ctx.fromData(data as TensorData, FP32::class) val x = ctx.fromFloatArray(Shape(256), FP32::class, FloatArray(256) { 0.125f }) - val got = ctx.ops.matmul(x, ctx.ops.transpose(w)).data.copyToFloatArray() + val got = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() var expect = 0f for (t in 0 until 256) expect += 0.125f * data.toFloatArray()[t] assertTrue(got[0].isFinite()); assertTrue(abs(got[0] - expect) < 1e-2f, "${got[0]} vs $expect") diff --git a/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedTransposeGoldenTest.kt b/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedTransposeGoldenTest.kt index a1596bdd6..fead65a6e 100644 --- a/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedTransposeGoldenTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedTransposeGoldenTest.kt @@ -23,7 +23,7 @@ import kotlin.test.assertTrue /** * #1034: a packed transpose is metadata. * - * `DefaultCpuOps.transpose` permutes the block grid byte by byte, because the packed matmul + * `relayoutPackedWeightForKernels` permutes the block grid byte by byte, because the packed matmul * kernels read their weight as input-block-major regardless of its declared shape — the contract * that made #968/#971 read garbage from a bare shape swap, and that #973 exists to write down. * A `TensorView` needs no such permutation: transposing swaps two strides and moves the block axis @@ -73,18 +73,18 @@ class PackedTransposeGoldenTest { } // 2. and it describes the same matrix as the physical block-grid permutation the kernels - // still need. `DefaultCpuOps.transpose` reorders the blocks to *input-block-major* — + // still need. `relayoutPackedWeightForKernels` reorders the blocks to *input-block-major* — // block (bI, o) at index `bI * rows + o` — because that is how the packed kernels read a // weight, whatever shape it declares (#968/#971; the contract #973 exists to write down). // So the permuted bytes are not the row-major encoding of the transposed matrix, and the // two paths are not interchangeable until #973 lands: decoded *as block-major*, they carry // exactly the values the zero-copy view exposes. - val physical = tensor.t() - assertEquals(Shape(shape[1], shape[0]), physical.shape, "${p.name}: ops.transpose shape") + val physical = ctx.ops.relayoutPackedWeightForKernels(tensor) + assertEquals(Shape(shape[1], shape[0]), physical.shape, "${p.name}: relayout shape") val permutedBytes = (physical.data as PackedBlockStorage).packedData assertTrue( permutedBytes.contentEquals(GoldenSupport.blockMajor(GoldenSupport.weightBlocks(p, ROWS, BLOCKS_PER_ROW, SEED))), - "${p.name}: ops.transpose must produce input-block-major bytes", + "${p.name}: the relayout must produce input-block-major bytes", ) val permuted = build(p, shape, permutedBytes) val block = FloatArray(p.blockSize) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/QuantizedMemSegMatmulTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/QuantizedMemSegMatmulTest.kt index aba082961..13909fcae 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/QuantizedMemSegMatmulTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/QuantizedMemSegMatmulTest.kt @@ -145,7 +145,7 @@ class QuantizedMemSegMatmulTest { @Suppress("UNCHECKED_CAST") val tensor: Tensor = VoidOpsTensor(q4k as TensorData, FP32::class) - val transposed = ops.transpose(tensor) + val transposed = ops.relayoutPackedWeightForKernels(tensor) assertEquals(Shape(Q4_KTensorData.BLOCK_SIZE, numBlocks), transposed.shape) assertTrue( transposed.data is Q4_KTensorData, @@ -179,7 +179,7 @@ class QuantizedMemSegMatmulTest { @Suppress("UNCHECKED_CAST") val tensor: Tensor = VoidOpsTensor(q6k as TensorData, FP32::class) - val transposed = ops.transpose(tensor) + val transposed = ops.relayoutPackedWeightForKernels(tensor) assertEquals(Shape(Q6_KTensorData.BLOCK_SIZE, numBlocks), transposed.shape) assertTrue( transposed.data is Q6_KTensorData, @@ -243,7 +243,7 @@ class QuantizedMemSegMatmulTest { sum } - val result = ops.matmul(input, ops.transpose(weight)) + val result = ops.matmulWeightTransposed(input, weight) val resultData = result.data.copyToFloatArray() assertEquals(outputDim, resultData.size) @@ -297,7 +297,7 @@ class QuantizedMemSegMatmulTest { sum } - val result = ops.matmul(input, ops.transpose(weight)) + val result = ops.matmulWeightTransposed(input, weight) val resultData = result.data.copyToFloatArray() assertEquals(outputDim, resultData.size) @@ -322,7 +322,7 @@ class QuantizedMemSegMatmulTest { val weight = q4Tensor(Shape(outputDim, inputDim), weightBytes, arena) val input = fpTensor(Shape(batchSize, inputDim), FloatArray(batchSize * inputDim) { 1f }) - val result = ops.matmul(input, ops.transpose(weight)) + val result = ops.matmulWeightTransposed(input, weight) assertEquals(Shape(batchSize, outputDim), result.shape) arena.close() } @@ -379,8 +379,8 @@ class QuantizedMemSegMatmulTest { ) val input: Tensor = VoidOpsTensor(inputData, FP32::class) - val transposedWeight = ops.transpose(weight) - assertTrue(transposedWeight.data is Q6_KTensorData, "transpose must preserve Q6_K packed layout") + val transposedWeight = ops.relayoutPackedWeightForKernels(weight) + assertTrue(transposedWeight.data is Q6_KTensorData, "the relayout must preserve Q6_K packed layout") val result = ops.matmul(input, transposedWeight) @@ -415,8 +415,8 @@ class QuantizedMemSegMatmulTest { Q4_KBlockTensorData(Shape(numBlocks, inputDim), weightBytes) as TensorData, FP32::class, ) - val transposedWeight = ops.transpose(weight) - assertTrue(transposedWeight.data is Q4_KTensorData, "transpose must preserve Q4_K packed layout") + val transposedWeight = ops.relayoutPackedWeightForKernels(weight) + assertTrue(transposedWeight.data is Q4_KTensorData, "the relayout must preserve Q4_K packed layout") // Rank 1, not rank 2 — a single-token hidden-state vector, not a `[1, in]` batch. val input = fpTensor(Shape(inputDim), FloatArray(inputDim) { (it + 1).toFloat() / inputDim }) diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeLazyTransposeGroundTruthReproTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeLazyTransposeGroundTruthReproTest.kt index 646f64e08..144e51d90 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeLazyTransposeGroundTruthReproTest.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeLazyTransposeGroundTruthReproTest.kt @@ -49,7 +49,7 @@ import sk.ainet.lang.types.FP32 * doc comment). That claim is only true if the bytes were ALREADY * kernel-native (input-block-major) before the swap. If a weight is * genuinely row-major (canonical bytes, `blocksPerInputDim > 1`), the lazy - * transpose does NOT reorder anything — so `ops.matmul(x, ops.transpose(w))` + * transpose does NOT reorder anything — so `ops.matmulWeightTransposed(x, w)` * on a canonically-packed weight hands the kernel bytes in the WRONG * physical order, silently, for every packed format with more than one * block per row. @@ -185,7 +185,7 @@ class NativeLazyTransposeGroundTruthReproTest { val ctxClassic = DirectCpuExecutionContext() val wClassic = ctxClassic.fromData(build(Shape(outputDim, inputDim), canonicalBytes), FP32::class) val xClassic = ctxClassic.fromFloatArray(Shape(1, inputDim), FP32::class, xf) - val yClassic = ctxClassic.ops.matmul(xClassic, ctxClassic.ops.transpose(wClassic)).data.copyToFloatArray() + val yClassic = ctxClassic.ops.matmulWeightTransposed(xClassic, wClassic).data.copyToFloatArray() val ctxPre = DirectCpuExecutionContext() val wPre = ctxPre.fromData(build(Shape(inputDim, outputDim), kernelNativeBytes), FP32::class) @@ -219,7 +219,7 @@ class NativeLazyTransposeGroundTruthReproTest { "groundTruth[0..3]=${yGroundTruth.take(4)} classic[0..3]=${yClassic.take(4)}", ) - // Regression contract (fixed): `ops.matmul(x, ops.transpose(w))` on a + // Regression contract (fixed): `ops.matmulWeightTransposed(x, w)` on a // canonically-packed weight — the "classic" path `linearProject` uses — // must match the SAME independent ground truth the "pre-transposed" // (skip-transpose, kernel-native-bytes) workaround already matched. diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index ccc457990..156d7e002 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -5882,6 +5882,7 @@ public final class sk/ainet/lang/tensor/ops/KspTensorOps : sk/ainet/lang/tensor/ public fun logSoftmax (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun lt (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; 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; 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; @@ -5892,6 +5893,7 @@ public final class sk/ainet/lang/tensor/ops/KspTensorOps : sk/ainet/lang/tensor/ public fun pow (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun powScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public fun rdivScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; public fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; @@ -5971,6 +5973,11 @@ public abstract interface class sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps public abstract fun convert (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/tensor/Tensor; } +public final class sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps$DefaultImpls { + public static fun matmulWeightTransposed (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; +} + public final class sk/ainet/lang/tensor/ops/MultiplyOperation : sk/ainet/lang/tensor/ops/BaseOperation { public fun ()V public fun (Ljava/util/Map;)V @@ -6146,6 +6153,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public static synthetic fun logSoftmax$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun lt (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public abstract 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; public abstract fun maxPool2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun maxPool2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun mean (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;)Lsk/ainet/lang/tensor/Tensor; @@ -6158,6 +6166,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public abstract fun pow (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun powScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public abstract fun rdivScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; public abstract fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; @@ -6201,8 +6210,10 @@ public final class sk/ainet/lang/tensor/ops/TensorOps$DefaultImpls { public static synthetic fun indexSelect$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun leakyRelu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun logSoftmax$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public static fun matmulWeightTransposed (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun maxPool2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun mean$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun scaledDotProductAttention$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun softmax$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun squeeze$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;ILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; @@ -6338,6 +6349,7 @@ public final class sk/ainet/lang/tensor/ops/VoidTensorOps : sk/ainet/lang/tensor public fun logSoftmax (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun lt (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; 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; 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; @@ -6348,6 +6360,7 @@ public final class sk/ainet/lang/tensor/ops/VoidTensorOps : sk/ainet/lang/tensor public fun pow (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun powScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public fun rdivScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; public fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt index 34154fb18..3f2098bf7 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt @@ -73,8 +73,10 @@ public open class Linear @kotlin.jvm.JvmOverloads constructor( val weight = params.weights().value val bias = params.biasOrNull()?.value - val weightTransposed = weight.t() - val matmulResult = input.matmul(weightTransposed) + // `x · Wᵀ` asked for directly (#973/#1096): the weight stays `[out, in]`, and a backend + // that would otherwise relayout a packed weight on every forward pass does it once. + // `matmul(x, weight.t())` is the same thing for dense data and a per-call copy for packed. + val matmulResult = ctx.ops.matmulWeightTransposed(input, weight) if (bias == null) return matmulResult // If input is a 1D vector, ensure bias is also 1D to avoid broadcasting to [1, out] diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index fe614093c..6c88be4d8 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -52,6 +52,44 @@ public interface TensorOps { @Diff @DarcValidated(by = "SKaiNET docs maintainers", on = "2026-05-24") public fun matmul(a: Tensor, b: Tensor): Tensor + + /** + * `x · Wᵀ` with the weight given as `[out, in]` — the primitive every reference implementation + * actually has (ggml's `mul_mat`, BLAS's `op(B)`), and the one a `Linear` layer wants + * (#973/#1096). + * + * `matmul(x, transpose(w))` is not the same thing for a **block-quantized** weight. Blocks + * quantize runs along the input dimension, so a real transpose would need requantization; what + * `transpose` does to a packed weight is a *layout conversion* wearing transpose's name, and + * since it happens per forward pass it copies the whole weight on every call. Asking for + * `x · Wᵀ` directly lets an implementation do the right thing — or nothing at all, when the + * weight is already in the order its kernel reads. + * + * The default keeps today's behaviour exactly, so no implementation has to change: it is + * literally `matmul(x, transpose(w))`. Backends override it where they can do better. + * + * @param x activations, `[batch, in]` or `[in]` + * @param weight `[out, in]` — **not** pre-transposed + */ + public fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor = + matmul(x, transpose(weight)) + + /** + * A block-packed weight's bytes rearranged into the order its kernels read (#973/#1096) — the + * operation `transpose` used to perform on packed data, under its own name. + * + * Callers who want a *product* should ask for [matmulWeightTransposed], which runs this once per + * weight rather than once per call. This exists for the ones who genuinely want the permuted + * bytes: a converter, a test, a benchmark. + * + * The default refuses, because a backend that has no packed kernels has no such order to + * produce. + */ + public fun relayoutPackedWeightForKernels(weight: Tensor): Tensor = + throw UnsupportedOperationException( + "this backend has no packed kernel order to relayout ${weight.data::class.simpleName} into", + ) + @Diff public fun transpose(tensor: Tensor): Tensor diff --git a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api index b0ea7ea57..c5646ee0e 100644 --- a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api +++ b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api @@ -34,6 +34,11 @@ public final class sk/ainet/lang/dag/GraphDslKt { public abstract interface class sk/ainet/lang/dag/GraphDslOps : sk/ainet/lang/tensor/ops/TensorOps { } +public final class sk/ainet/lang/dag/GraphDslOps$DefaultImpls { + public static fun matmulWeightTransposed (Lsk/ainet/lang/dag/GraphDslOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/dag/GraphDslOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; +} + public final class sk/ainet/lang/dag/GraphDslOpsGraphDslKt { public static final fun abs (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun abs$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue;