From e80aee485ec6a95e5f430c8732ed0e2b52832d54 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 22:25:00 +0200 Subject: [PATCH] feat(ops)!: matmulWeightTransposed as the primitive; transpose refuses a packed weight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1096, the last slice of #973. `ops.transpose` on block-quantized data was never a transpose. Blocks quantize runs along the input dimension, so a real one needs requantization; what happened instead was a layout conversion wearing transpose's name, with a shape label that lied about the bytes. It was not an involution — `transpose(transpose(W))` is a different matrix for a non-square block grid — and since #969 it copied the whole weight on *every* call, which for `Linear.onForward` meant every forward pass. - `TensorOps.matmulWeightTransposed(x, weight)`: `x · Wᵀ` with the weight as `[out, in]` — the primitive ggml (`mul_mat`) and BLAS (`op(B)`) actually have. Defaulted to `matmul(x, transpose(w))`, so no implementation has to change; `DefaultCpuOps` overrides it to relayout a packed weight **once per weight** and reuse it, which is the point. - `TensorOps.relayoutPackedWeightForKernels(weight)`: the conversion under its own name, for callers who want the permuted bytes rather than a product. Defaulted to refusing. - `DefaultCpuOps.transpose` **throws** for a heap packed weight, naming the primitive, the explicit relayout, and why the operation is not definable. Scoped to the heap tier: the MemorySegment kernels read canonical bytes, so a shape swap is genuinely correct there — census contradiction #3, now stated in code rather than implied by a marker interface. Narrow floats keep their view rewrap (#888). - `Linear.onForward` asks for the product instead of transposing. Breaking, deliberately: 27 in-repo call sites moved to one of the two new entry points. The user has confirmed the downstream break is acceptable and that versions are pinned. The relayout cache is bounded (64 weights) and keyed by the identity of the packed bytes: a model's weights are stable, so the first forward converts and every later one reuses; beyond the bound the oldest is dropped and reconverted, which is exactly the old per-call behaviour. Tests: the primitive computes exactly what relayout-then-matmul did; repeated calls keep agreeing with the explicit relayout; transpose refuses with a message naming the primitive, the relayout and the reason; a dense weight transposes as it always did. Gate: scripts/pr-gate.sh — all legs passed; --golden passed. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-backend-cpu.api | 2 + .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 147 ++++++++++++++---- .../golden/PackedMatmulDispatchParityTest.kt | 4 +- .../tensor/ops/MatmulWeightTransposedTest.kt | 95 +++++++++++ .../tensor/ops/PackedMatmulDispatchTest.kt | 8 +- .../tensor/ops/RegistryMatmulDispatchTest.kt | 8 +- .../exec/golden/PackedTransposeGoldenTest.kt | 10 +- .../tensor/ops/QuantizedMemSegMatmulTest.kt | 18 +-- ...NativeLazyTransposeGroundTruthReproTest.kt | 6 +- .../api/jvm/skainet-lang-core.api | 13 ++ .../kotlin/sk/ainet/lang/nn/Linear.kt | 6 +- .../sk/ainet/lang/tensor/ops/TensorOps.kt | 38 +++++ .../api/jvm/skainet-lang-dag.api | 5 + 13 files changed, 297 insertions(+), 63 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/MatmulWeightTransposedTest.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 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;