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 {