From 1ef680759273ce6fe4e8d80069a95481e57b5ff2 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 25 Aug 2026 17:28:51 +0200 Subject: [PATCH] fix(cpu): stop relayouting a packed weight for kernels this backend does not have MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1124. The common DefaultCpuOps computed packed matmul wrongly whenever the KernelRegistry was empty — silently, by several times the magnitude of the answer. Ground truth is the weight's own toFloatArray(), so the packed path was disagreeing with its own decoder. matmulWeightTransposed relayouted the weight into the order the *JVM* vectorized kernels read, then handed the result to whatever computed the product. On the JVM that is chooseQuantizedMatmul, which addresses packedData in feed order deliberately and is correct. The common implementation has no such kernel: it decodes through views, and what it was decoding was a tensor whose shape says [in, out] while its blocks still run along the original input dimension. That combination is not canonical and is not input-block-major-at-that-shape; it is a private artifact that only its producer understands, and reading it as either gives plausible garbage. So the common path no longer relayouts. KernelDispatch.matmul already wants the weight output-major, which is exactly the shape an [out, in] weight has, so the canonical view goes in with nothing rearranged. DefaultCpuOpsJvm overrides matmulWeightTransposed to keep the relayout-and-cache path, where the kernels that consume those bytes actually exist and #1096's once-per-weight conversion still pays. Also plumbs block order through to views: PackedBlockStorage.blockOrder, defaulted to ROW_MAJOR so nothing existing changes, forwarded by packedView into Layout.blocked — which has accepted a blockOrder since #1094 and never received one. That closes the information gap #1120 needs. It is worth recording what it cannot do: the relayout artifact above has no honest block order to declare, so marking it was the first fix I tried and it produced zeros. Why no test caught the original bug: skainet-backend-cpu's JVM tests always resolve to DefaultCpuOpsJvm, and its native/JS/Wasm tests do use the common implementation but their platform factories register ScalarKernelProvider first. Every packed-matmul test in the tree, including those added for #973, #1096 and #1108, validated a configuration that was not the broken one. The new test pins the empty-registry configuration explicitly, in commonTest so it runs where that path is real. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-backend-cpu.api | 1 + .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 55 +++++++-- .../ops/PackedMatmulEmptyRegistryTest.kt | 105 ++++++++++++++++++ .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 33 ++++++ .../api/jvm/skainet-lang-core.api | 35 ++++-- .../kotlin/sk/ainet/lang/memory/TensorView.kt | 19 +++- .../ainet/lang/tensor/data/Q4_0TensorData.kt | 7 +- .../ainet/lang/tensor/data/Q4_KTensorData.kt | 7 +- .../ainet/lang/tensor/data/Q5_0TensorData.kt | 5 + .../ainet/lang/tensor/data/Q5_1TensorData.kt | 5 + .../ainet/lang/tensor/data/Q5_KTensorData.kt | 7 +- .../ainet/lang/tensor/data/Q6_KTensorData.kt | 7 +- .../ainet/lang/tensor/data/Q8_0TensorData.kt | 7 +- .../lang/tensor/storage/PackedBlockStorage.kt | 21 ++++ 14 files changed, 290 insertions(+), 24 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulEmptyRegistryTest.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 a52798738..6757ae6ee 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 @@ -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; 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 2e9ab83fd..d19d8e384 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 @@ -918,14 +918,53 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory @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) + // 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 matmulWeightTransposedViaViews(x: Tensor, weight: Tensor): Tensor? { + 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(outShape, x.dtype, outArray) as sk.ainet.lang.tensor.data.TensorData + return newTensor(outData, x.dtype, x, weight) } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulEmptyRegistryTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PackedMatmulEmptyRegistryTest.kt new file mode 100644 index 000000000..eb38cb328 --- /dev/null +++ b/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 TensorData>> = listOf( + Triple("Q8_0", 34) { s, b -> Q8_0BlockTensorData(s, b) as TensorData }, + Triple("Q4_0", 18) { s, b -> Q4_0BlockTensorData(s, b) as TensorData }, + Triple("Q5_0", 22) { s, b -> Q5_0BlockTensorData(s, b) as TensorData }, + Triple("Q5_1", 24) { s, b -> Q5_1BlockTensorData(s, b) as TensorData }, + ) + // 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 = 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(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]}", + ) + } + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 00d149c89..0f5992404 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -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 matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { + 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) + } + + /** Relayouted weights keyed by the identity of the bytes they came from (#1096). */ + private val prepackedWeightsJvm: MutableList>> = 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 matmul(a: Tensor, b: Tensor): Tensor { // `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 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 869f5e484..5c1384f61 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 @@ -1265,8 +1265,8 @@ public final class sk/ainet/lang/memory/TensorView { public final class sk/ainet/lang/memory/TensorView$Companion { public final fun dense (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; public static synthetic fun dense$default (Lsk/ainet/lang/memory/TensorView$Companion;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; - public final fun packed (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/memory/TensorView; - public static synthetic fun packed$default (Lsk/ainet/lang/memory/TensorView$Companion;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; + public final fun packed (Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/BlockOrder;)Lsk/ainet/lang/memory/TensorView; + public static synthetic fun packed$default (Lsk/ainet/lang/memory/TensorView$Companion;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/BlockOrder;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView; } public final class sk/ainet/lang/memory/TernaryBlockDecoder : sk/ainet/lang/memory/BlockDecoder { @@ -4890,12 +4890,14 @@ public final class sk/ainet/lang/tensor/data/Q4MemorySegmentTensorData$Companion public final class sk/ainet/lang/tensor/data/Q4_0BlockTensorData : sk/ainet/lang/tensor/data/Q4_0TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q4_0BlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; public synthetic fun get ([I)Ljava/lang/Object; public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockScale (I)F public fun getBlockSize ()I public fun getCode (II)B @@ -4950,7 +4952,8 @@ public final class sk/ainet/lang/tensor/data/Q4_0TensorDataKt { public final class sk/ainet/lang/tensor/data/Q4_KBlockTensorData : sk/ainet/lang/tensor/data/Q4_KTensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q4_KBlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; @@ -4958,6 +4961,7 @@ public final class sk/ainet/lang/tensor/data/Q4_KBlockTensorData : sk/ainet/lang public fun getBlockCount ()I public fun getBlockD (I)F public fun getBlockDMin (I)F + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockSize ()I public fun getCode (II)I public fun getElementCount ()J @@ -5014,12 +5018,14 @@ public final class sk/ainet/lang/tensor/data/Q4_KTensorDataKt { public final class sk/ainet/lang/tensor/data/Q5_0BlockTensorData : sk/ainet/lang/tensor/data/Q5_0TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q5_0BlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; public synthetic fun get ([I)Ljava/lang/Object; public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockSize ()I public fun getElementCount ()J public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; @@ -5060,12 +5066,14 @@ public final class sk/ainet/lang/tensor/data/Q5_0TensorData$DefaultImpls { public final class sk/ainet/lang/tensor/data/Q5_1BlockTensorData : sk/ainet/lang/tensor/data/Q5_1TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q5_1BlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; public synthetic fun get ([I)Ljava/lang/Object; public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockSize ()I public fun getElementCount ()J public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; @@ -5106,7 +5114,8 @@ public final class sk/ainet/lang/tensor/data/Q5_1TensorData$DefaultImpls { public final class sk/ainet/lang/tensor/data/Q5_KBlockTensorData : sk/ainet/lang/tensor/data/Q5_KTensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q5_KBlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; @@ -5114,6 +5123,7 @@ public final class sk/ainet/lang/tensor/data/Q5_KBlockTensorData : sk/ainet/lang public fun getBlockCount ()I public fun getBlockD (I)F public fun getBlockDMin (I)F + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockSize ()I public fun getCode (II)I public fun getElementCount ()J @@ -5174,13 +5184,15 @@ public final class sk/ainet/lang/tensor/data/Q5_KTensorDataKt { public final class sk/ainet/lang/tensor/data/Q6_KBlockTensorData : sk/ainet/lang/tensor/data/Q6_KTensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q6_KBlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; public synthetic fun get ([I)Ljava/lang/Object; public fun getBlockCount ()I public fun getBlockD (I)F + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockSize ()I public fun getCode (II)I public fun getElementCount ()J @@ -5264,12 +5276,14 @@ public final class sk/ainet/lang/tensor/data/Q8MemorySegmentTensorData$Companion public final class sk/ainet/lang/tensor/data/Q8_0BlockTensorData : sk/ainet/lang/tensor/data/Q8_0TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { public static final field Companion Lsk/ainet/lang/tensor/data/Q8_0BlockTensorData$Companion; - public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;)V + public synthetic fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/memory/BlockOrder;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun copyToFloatArray ()[F public fun dequantizeBlock (I[FI)V public fun get ([I)Ljava/lang/Byte; public synthetic fun get ([I)Ljava/lang/Object; public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockScale (I)F public fun getBlockSize ()I public fun getCode (II)B @@ -5373,6 +5387,7 @@ public final class sk/ainet/lang/tensor/data/Ternary2BitTensorData : sk/ainet/la public fun get ([I)Ljava/lang/Byte; public synthetic fun get ([I)Ljava/lang/Object; public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public fun getBlockSize ()I public fun getElementCount ()J public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; @@ -7166,6 +7181,7 @@ public abstract interface class sk/ainet/lang/tensor/storage/PackedBlockStorage public abstract fun dequantizeBlock (I[FI)V public static synthetic fun dequantizeBlock$default (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;I[FIILjava/lang/Object;)V public abstract fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; public abstract fun getBlockSize ()I public fun getElementCount ()J public abstract fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; @@ -7182,6 +7198,7 @@ public abstract interface class sk/ainet/lang/tensor/storage/PackedBlockStorage public final class sk/ainet/lang/tensor/storage/PackedBlockStorage$DefaultImpls { public static synthetic fun dequantizeBlock$default (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;I[FIILjava/lang/Object;)V + public static fun getBlockOrder (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)Lsk/ainet/lang/memory/BlockOrder; public static fun getElementCount (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)J public static fun getPackedView (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)Lsk/ainet/lang/memory/TensorView; public static fun getPhysicalBytes (Lsk/ainet/lang/tensor/storage/PackedBlockStorage;)J diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt index 3369963c9..ccebc84fa 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/TensorView.kt @@ -299,8 +299,23 @@ public class TensorView( * blocks in row-major block order, decoded by [decoder]. Slicing and transposing address * whole blocks — the bytes are never touched. */ - public fun packed(storage: Storage, shape: Shape, encoding: TensorEncoding, decoder: BlockDecoder, dtype: DType = FP32, id: TensorId? = null): TensorView = - TensorView(shape, Format(dtype, encoding), Layout.blocked(shape, decoder.blockSize, decoder.bytesPerBlock), storage, id, decoder) + public fun packed( + storage: Storage, + shape: Shape, + encoding: TensorEncoding, + decoder: BlockDecoder, + dtype: DType = FP32, + id: TensorId? = null, + blockOrder: BlockOrder = BlockOrder.ROW_MAJOR, + ): TensorView = + TensorView( + shape, + Format(dtype, encoding), + Layout.blocked(shape, decoder.blockSize, decoder.bytesPerBlock, blockOrder = blockOrder), + storage, + id, + decoder, + ) } } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_0TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_0TensorData.kt index fdeca6345..9aaa4f720 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_0TensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_0TensorData.kt @@ -64,7 +64,12 @@ public interface Q4_0TensorData : TensorData { */ public class Q4_0BlockTensorData( initialShape: Shape, - private val data: ByteArray + private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q4_0TensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_KTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_KTensorData.kt index 8a0655234..cf3a24b9d 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_KTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q4_KTensorData.kt @@ -91,7 +91,12 @@ public interface Q4_KTensorData : TensorData { */ public class Q4_KBlockTensorData( initialShape: Shape, - private val data: ByteArray + private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q4_KTensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_0TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_0TensorData.kt index 923a4c58e..8a1cc1ad1 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_0TensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_0TensorData.kt @@ -38,6 +38,11 @@ public interface Q5_0TensorData : TensorData { public class Q5_0BlockTensorData( initialShape: Shape, private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q5_0TensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_1TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_1TensorData.kt index 6f92e8812..25ec6b92e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_1TensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_1TensorData.kt @@ -43,6 +43,11 @@ public interface Q5_1TensorData : TensorData { public class Q5_1BlockTensorData( initialShape: Shape, private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q5_1TensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_KTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_KTensorData.kt index c94bbd927..51b25d3cc 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_KTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q5_KTensorData.kt @@ -97,7 +97,12 @@ public interface Q5_KTensorData : TensorData { */ public class Q5_KBlockTensorData( initialShape: Shape, - private val data: ByteArray + private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q5_KTensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q6_KTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q6_KTensorData.kt index 90f8af1db..c6712d22c 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q6_KTensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q6_KTensorData.kt @@ -82,7 +82,12 @@ public interface Q6_KTensorData : TensorData { */ public class Q6_KBlockTensorData( initialShape: Shape, - private val data: ByteArray + private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q6_KTensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0TensorData.kt index 5ea59a5dd..666f5c27a 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0TensorData.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Q8_0TensorData.kt @@ -51,7 +51,12 @@ public interface Q8_0TensorData : TensorData { */ public class Q8_0BlockTensorData( initialShape: Shape, - private val data: ByteArray + private val data: ByteArray, + /** + * Which order [data]'s blocks are physically in (#1120/#1124). `ROW_MAJOR` — canonical, as + * every GGUF-shaped producer writes — unless this weight was relayouted for the packed kernels. + */ + override val blockOrder: sk.ainet.lang.memory.BlockOrder = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR, ) : Q8_0TensorData, PackedBlockStorage { /** The façade over the packed bytes (SKEEP-003 §4.1): see [PackedBlockStorage.packedView]. */ diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt index 7501a8765..6cefa7266 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt @@ -32,6 +32,26 @@ public interface PackedBlockStorage { /** Raw packed byte data containing all blocks. */ public val packedData: ByteArray + /** + * The order [packedData]'s blocks are physically in (#1120, #1124). + * + * Canonical storage — every GGUF-shaped producer — is + * [sk.ainet.lang.memory.BlockOrder.ROW_MAJOR]: block `(o, b)` at flat index `o * blocksPerRow + b`. + * A weight that has been relayouted for the packed matmul kernels is + * [sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR] instead, and the two coincide only at one + * block per row. + * + * Before this existed, a relayouted weight was wrapped in a data type that still claimed to be + * canonical. The kernels that read [packedData] directly were unaffected — they were addressing + * it in feed order deliberately — but anything decoding through [packedView] read the wrong + * blocks and returned plausible garbage (#1124, and #973/#968 before it). Declaring the order is + * what lets both readers be right about the same bytes. + * + * Defaults to `ROW_MAJOR`, so every existing implementation is unchanged. + */ + public val blockOrder: sk.ainet.lang.memory.BlockOrder + get() = sk.ainet.lang.memory.BlockOrder.ROW_MAJOR + /** Physical byte size of the packed data. */ public val physicalBytes: Long get() = packedData.size.toLong() @@ -68,6 +88,7 @@ public interface PackedBlockStorage { shape = shape, encoding = encoding, decoder = sk.ainet.lang.memory.PackedBlockDecoder(this), + blockOrder = blockOrder, ) public fun toFloatArray(): FloatArray {