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 7a5ccdd43..fc9c123df 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 @@ -458,6 +458,67 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory private val q5_1Kernel by lazy { resolveProvider { it.matmulQ5_1() != null }?.matmulQ5_1() } private val q5_0Kernel by lazy { resolveProvider { it.matmulQ5_0() != null }?.matmulQ5_0() } + /** + * `cols / blockSize`, i.e. the number of quantization blocks per output row — + * validated so a misaligned packed tensor fails loudly here rather than + * silently truncating a partial trailing block during [transposePackedBlocks]. + */ + private fun requirePackedBlockAligned(cols: Int, blockSize: Int, formatName: String): Int { + require(cols % blockSize == 0) { + "$formatName transpose: inputDim $cols is not a multiple of the block size $blockSize " + + "— packed weight is not row-block-aligned, cannot compute a physical block-grid transpose" + } + return cols / blockSize + } + + /** + * Physically transposes a packed weight's block grid from *row-major* + * (canonical) order into the *input-block-major* order the packed-quant + * matmul kernels require. + * + * Canonical storage — what a fresh `[outputDim, inputDim]`-shaped packed + * tensor has, whether loaded verbatim from GGUF or built any other + * row-major way — groups blocks per output row: for row `o`, its + * `blocksPerInputDim` blocks are contiguous, then row `o + 1`. Flat block + * index = `o * blocksPerInputDim + blockIdx`. + * + * Every packed-quant native/Panama/scalar matmul kernel instead reads + * `weight + (blockIdx * outputDim + o) * bytesPerBlock` — see e.g. the + * `Per-block packed weight layout` header comment in `q5_0_matmul.c` / + * `q4k_matmul.c` / etc. — because for a FIXED input block, all `outputDim` + * rows' corresponding block bytes are consecutive, letting the kernel's + * block-outer/row-inner loop read weight memory sequentially instead of + * striding `outputDim * bytesPerBlock` per input block. + * + * These two orderings are literal transposes of the `(outputDim, + * blocksPerInputDim)` block grid (treating each `bytesPerBlock`-sized + * block as an atomic item) and coincide only when `blocksPerInputDim == 1`. + * For every wider weight — i.e. essentially every real model, since + * `inputDim` is almost always many multiples of the 32/256-element block + * size — reordering the *shape* without reordering the *bytes* hands the + * kernel physically wrong data: it silently reads block `bI` of row `o` + * from where block `o`'s row `bI`-th chunk actually lives. This is the + * root cause of the SKaiNET-transformers#307 all-zero Q5_0/Q5_1 matmul + * report (general to every packed format the native tier serves, not + * Q5-specific — see `NativeLazyTransposeGroundTruthReproTest`). + */ + private fun transposePackedBlocks( + packed: ByteArray, + outputDim: Int, + blocksPerInputDim: Int, + bytesPerBlock: Int, + ): ByteArray { + val out = ByteArray(packed.size) + for (o in 0 until outputDim) { + for (blockIdx in 0 until blocksPerInputDim) { + val srcOff = (o * blocksPerInputDim + blockIdx) * bytesPerBlock + val dstOff = (blockIdx * outputDim + o) * bytesPerBlock + packed.copyInto(out, dstOff, srcOff, srcOff + bytesPerBlock) + } + } + return out + } + /** * Platform-neutral packed-quant matmul: `FP32 input × packed-quant weight`, * resolving the kernel via [KernelRegistry] (scalar on Native/JS/WASM, Panama/ @@ -714,28 +775,68 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory val rows = tensor.shape[rank - 2] val cols = tensor.shape[rank - 1] - // Lazy transpose for heap-packed quant weights (Q4_K/Q6_K/Q5_1/Q5_0): the - // matmul kernels index the packed bytes input-block-major from the post-swap - // (inputDim, outputDim), so transpose is a pure shape swap — same bytes, no copy. - // Lets `ops.matmul(x, ops.transpose(W))` run on every platform without a dequant - // round-trip. (The JVM ops intercept Q4_K/Q6_K + MemSeg before reaching here.) + // 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 -> return newTensor(Q4_KBlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) - is Q5_KTensorData -> return newTensor(Q5_KBlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) - is Q6_KTensorData -> return newTensor(Q6_KBlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) - is Q5_1TensorData -> return newTensor(Q5_1BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) - is Q5_0TensorData -> return newTensor(Q5_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) - // Q8_0 / Q4_0 lazy transpose: rewrap the same input-block-major bytes - // with flipped shape (bytes are layout-agnostic to the [out,in] kernel - // convention) so a packed weight (e.g. gemma's tied Q8_0 lm_head) - // survives linearProject's transpose instead of hitting the generic - // FP32 path (Byte→Float ClassCastException). This `when` now covers - // every quant type chooseQuantizedMatmulHeap dispatches — i.e. every - // packed type that can be a matmul weight. See transformers #178. - is Q8_0TensorData -> return newTensor(Q8_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) - is Q4_0TensorData -> return newTensor(Q4_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) + is Q4_KTensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q4_KTensorData.BLOCK_SIZE, "Q4_K") + val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q4_KTensorData.BYTES_PER_BLOCK) + return newTensor(Q4_KBlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) + } + is Q5_KTensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q5_KTensorData.BLOCK_SIZE, "Q5_K") + val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q5_KTensorData.BYTES_PER_BLOCK) + return newTensor(Q5_KBlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) + } + is Q6_KTensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q6_KTensorData.BLOCK_SIZE, "Q6_K") + val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q6_KTensorData.BYTES_PER_BLOCK) + return newTensor(Q6_KBlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) + } + is Q5_1TensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q5_1TensorData.BLOCK_SIZE, "Q5_1") + val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q5_1TensorData.BYTES_PER_BLOCK) + return newTensor(Q5_1BlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) + } + is Q5_0TensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q5_0TensorData.BLOCK_SIZE, "Q5_0") + val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q5_0TensorData.BYTES_PER_BLOCK) + return newTensor(Q5_0BlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) + } + // Q8_0 / Q4_0: same physical block-grid transpose as the arms above — + // this `when` covers every quant type chooseQuantizedMatmulHeap + // dispatches, i.e. every packed type that can be a matmul weight + // (originally retrofitted for the Byte→Float ClassCastException gap, + // see transformers #178; the shape-swap-only version of these arms + // carried the same block-order bug as Q5_0/Q5_1 above). + is Q8_0TensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q8_0TensorData.BLOCK_SIZE, "Q8_0") + val reordered = transposePackedBlocks(d.packedData, rows, blocksPerInputDim, Q8_0TensorData.BYTES_PER_BLOCK) + return newTensor(Q8_0BlockTensorData(Shape(cols, rows), reordered) as TensorData, tensor.dtype, tensor) + } + is Q4_0TensorData -> { + val blocksPerInputDim = requirePackedBlockAligned(cols, Q4_0TensorData.BLOCK_SIZE, "Q4_0") + 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 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 6593645b7..5a1923e45 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 @@ -35,11 +35,22 @@ class PackedMatmulDispatchTest { } private fun le16(b: ByteArray, o: Int, h: Int) { b[o] = (h and 0xFF).toByte(); b[o + 1] = ((h ushr 8) and 0xFF).toByte() } - /** Random block-major Q5_1 bytes for [out,in] + the FP32 weight they dequantize to (row-major). */ + /** + * Random CANONICAL (row-major: for output row `o`, its `blocks` input-blocks + * are contiguous, then row `o + 1`) Q5_1 bytes for `[out,in]` + the FP32 + * weight they dequantize to (row-major) — i.e. bytes shaped the way a + * `[outputDim, inputDim]`-shaped packed tensor actually has them (verbatim + * GGUF load, or any other row-major producer), NOT the kernel-native + * `(blockIdx * outputDim + o)` layout the matmul kernels want. `ops + * .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. + */ 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) for (o in 0 until outDim) for (bI in 0 until blocks) { - val off = (bI * outDim + o) * 24; val dst = o * inDim + bI * 32 + val off = (o * blocks + bI) * 24; val dst = o * inDim + bI * 32 val d = rng.nextFloat() * 0.05f + 0.01f; val m = rng.nextFloat() - 0.5f le16(bytes, off, half(d)); le16(bytes, off + 2, half(m)) val qh = IntArray(4) { rng.nextInt(256) }; for (k in 0 until 4) bytes[off + 4 + k] = qh[k].toByte() @@ -53,11 +64,11 @@ class PackedMatmulDispatchTest { return bytes to wf } - /** Random block-major Q4_K bytes for [out,in] + the FP32 weight. */ + /** Random CANONICAL (row-major, see [q5_1]) Q4_K bytes for [out,in] + the FP32 weight. */ private fun q4_k(inDim: Int, outDim: Int, rng: Random): Pair { val blocks = inDim / 256; val bytes = ByteArray(outDim * blocks * 144); val wf = FloatArray(outDim * inDim) for (o in 0 until outDim) for (bI in 0 until blocks) { - val off = (bI * outDim + o) * 144; val dst = o * inDim + bI * 256 + val off = (o * blocks + bI) * 144; val dst = o * inDim + bI * 256 val d = rng.nextFloat() * 0.02f + 0.005f; val dMin = rng.nextFloat() * 0.02f + 0.005f le16(bytes, off, half(d)); le16(bytes, off + 2, half(dMin)) for (k in 0 until 140) bytes[off + 4 + k] = rng.nextInt(256).toByte() @@ -80,11 +91,11 @@ class PackedMatmulDispatchTest { return bytes to wf } - /** Random block-major Q6_K bytes for [out,in] + the FP32 weight. */ + /** Random CANONICAL (row-major, see [q5_1]) Q6_K bytes for [out,in] + the FP32 weight. */ private fun q6_k(inDim: Int, outDim: Int, rng: Random): Pair { val blocks = inDim / 256; val bytes = ByteArray(outDim * blocks * 210); val wf = FloatArray(outDim * inDim) for (o in 0 until outDim) for (bI in 0 until blocks) { - val off = (bI * outDim + o) * 210; val dst = o * inDim + bI * 256 + val off = (o * blocks + bI) * 210; val dst = o * inDim + bI * 256 for (k in 0 until 208) bytes[off + k] = rng.nextInt(256).toByte() val d = rng.nextFloat() * 0.01f + 0.002f; le16(bytes, off + 208, half(d)) for (h in 0..1) { 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 4364b89b6..ee43f3576 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 @@ -32,7 +32,6 @@ import sk.ainet.lang.tensor.data.Q4_0TensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.Q8MemorySegmentMarker import sk.ainet.lang.tensor.data.Q8MemorySegmentTensorData -import sk.ainet.lang.tensor.data.Q4_KBlockTensorData import sk.ainet.lang.tensor.data.Q4_KTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.BF16 @@ -225,24 +224,22 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") return newTensor(transposed as TensorData, tensor.dtype, tensor) } - // Lazy transpose for Q4_K packed data: swap shape, keep the packed byte - // array untouched. `JvmQuantizedVectorKernels.matmulQ4_KVec` derives its - // byte offsets from (inputDim, outputDim) via `(blockIdx * outputDim + o)` - // and the packed layout is input-block-major (all output rows for a given - // input block packed contiguously), so the same bytes produce the right - // values under the swapped shape. This is the DSL-path counterpart of the - // Q4/Q8 MemSeg lazy transpose: Q4_K weights can flow through - // `ops.matmul(x, ops.transpose(W))` without a dequant round-trip. - if (data is Q4_KTensorData) { - val packedData = data.packedData - val transposed = Q4_KBlockTensorData(Shape(cols, rows), packedData) - @Suppress("UNCHECKED_CAST") - return newTensor(transposed as TensorData, tensor.dtype, tensor) - } + // Q4_K (ByteArray-backed `Q4_KBlockTensorData`) transpose used to be + // intercepted here with a bare shape swap, on the claim that + // `JvmQuantizedVectorKernels.matmulQ4_KVec`'s `(blockIdx * outputDim + o)` + // byte addressing matches the packed layout "as-is". That's only true + // when there's a single quant block per row; for any wider row it silently + // fed the kernel physically wrong bytes — the same root cause as the + // Q5_0/Q5_1/Q8_0/Q4_0/Q5_K/Q6_K bug fixed in `DefaultCpuOpsBase.transpose` + // (see `transposePackedBlocks`, SKaiNET-transformers#307). Falling through + // to `super.transpose()` picks up the corrected, physically-reordering + // implementation instead of duplicating (and re-diverging from) it here. + // // Narrow-float input-major lazy transpose is handled in DefaultCpuOpsBase too — // nothing above intercepts it, so it falls through. Issue #888. - // Q6_K / Q5_1 / Q5_0 lazy transpose is handled in DefaultCpuOpsBase - // (block-major, shared with Native); the JVM ops don't intercept them here. + // Q6_K / Q5_1 / Q5_0 / Q4_K / Q8_0 / Q4_0 transpose is handled in + // DefaultCpuOpsBase (block-major, shared with Native); the JVM ops don't + // intercept them here. // MemorySegment FP32 fast path: physical transpose via SIMD. // Uses Arena.ofAuto() so the result segment is reclaimed by GC // when the wrapping Tensor is no longer reachable. Earlier 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 30c3fd074..8cd194fc4 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 @@ -150,14 +150,20 @@ class QuantizedMemSegMatmulTest { transposed.data is Q4_KTensorData, "transpose must preserve Q4_K packed layout, got ${transposed.data::class.simpleName}" ) - // Lazy invariant: the packed byte array must be the SAME reference — - // no copy, no re-layout. This is what distinguishes the specialized - // branch from the fallback per-element transpose (which would crash - // on Byte → Float casts, the very regression this test guards). + // `cols == BLOCK_SIZE` means exactly one block per row (blocksPerInputDim == + // 1), so the physical block-grid transpose (see `DefaultCpuOpsBase + // .transposePackedBlocks`) is a content-preserving permutation here — but it + // is still a genuine copy (a new array), not the old bare shape-swap-only + // "zero-copy" behaviour, which was only byte-order-correct in this exact + // single-block-per-row case and silently wrong for every wider row + // (SKaiNET-transformers#307). Assert content equality, not reference + // identity — the specialized branch (vs. the fallback per-element + // transpose, which would crash on Byte → Float casts, the original + // regression this test guarded) is still exercised either way. val transposedPacked = (transposed.data as Q4_KTensorData).packedData assertTrue( - transposedPacked === bytes, - "Q4_K lazy transpose must keep the same packedData reference (zero-copy)" + transposedPacked.contentEquals(bytes), + "Q4_K transpose must preserve block content for a single-block-per-row weight", ) } @@ -178,10 +184,12 @@ class QuantizedMemSegMatmulTest { transposed.data is Q6_KTensorData, "transpose must preserve Q6_K packed layout, got ${transposed.data::class.simpleName}" ) + // See the Q4_K case above: single block per row (cols == BLOCK_SIZE) makes + // the physical block-grid transpose content-preserving but still a copy. val transposedPacked = (transposed.data as Q6_KTensorData).packedData assertTrue( - transposedPacked === bytes, - "Q6_K lazy transpose must keep the same packedData reference (zero-copy)" + transposedPacked.contentEquals(bytes), + "Q6_K transpose must preserve block content for a single-block-per-row weight", ) } 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 new file mode 100644 index 000000000..646f64e08 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeLazyTransposeGroundTruthReproTest.kt @@ -0,0 +1,297 @@ +package sk.ainet.exec.kernel + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertTrue +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Q4_0BlockTensorData +import sk.ainet.lang.tensor.data.Q4_KBlockTensorData +import sk.ainet.lang.tensor.data.Q5_0BlockTensorData +import sk.ainet.lang.tensor.data.Q5_1BlockTensorData +import sk.ainet.lang.tensor.data.Q5_KBlockTensorData +import sk.ainet.lang.tensor.data.Q6_KBlockTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 + +/** + * Ground-truth variant of the repro: unlike [NativeLazyTransposeAllZeroReproTest] + * (which feeds the SAME flat byte array into both the "classic" and + * "pre-transposed" constructions — which is mathematically guaranteed to + * produce identical output no matter what, since `ops.transpose` is proven- + * by-inspection to be a pure metadata relabel over the same bytes), this test + * builds two bytewise-DIFFERENT packings of the SAME logical weight matrix: + * + * - "canonical" bytes: true row-major — for output row `o`, its + * `blocksPerInputDim` blocks are stored contiguously, then row `o+1`. This + * is what a loader constructing a `[outputDim, inputDim]`-shaped weight + * tensor "the way GGUF stores a row-major matrix" would naturally produce, + * and it is what `PackedBlockStorage.toFloatArray()` (block-sequential + * dequant) reconstructs as ground truth — reshaping its output as + * `[outputDim][inputDim]` gives the authoritative W[o][i] used nowhere near + * the matmul kernel under test. + * - "kernel-native" bytes: input-block-major — `(blockIdx * outputDim + o)` — + * the physical order every native matmul kernel here actually assumes + * (see the q5_0_matmul.c / q4_0_matmul.c header comments). Same logical + * blocks, same per-(o,blockIdx) content, just placed at different flat + * offsets. + * + * `ops.transpose` claims a packed weight can flow from a `[outputDim, + * inputDim]`-shaped tensor to `[inputDim, outputDim]` for free — same bytes, + * just a shape relabel — because "the matmul kernels index the packed bytes + * input-block-major from the post-swap shape" (`DefaultCpuOps.transpose` + * 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))` + * 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. + */ +class NativeLazyTransposeGroundTruthReproTest { + + @BeforeTest + fun forceNativeProviderOnly() { + KernelRegistry.clearForTesting() + KernelRegistry.register(NativeKernelProvider) + assertTrue(NativeKernelProvider.isAvailable(), "native provider must be available for this repro") + } + + @AfterTest + fun reset() = KernelRegistry.clearForTesting() + + private fun le16(b: ByteArray, o: Int, h: Int) { + b[o] = (h and 0xFF).toByte(); b[o + 1] = ((h ushr 8) and 0xFF).toByte() + } + + /** One synthetic Q5_0/Q4_0/Q8_0/Q5_1 block's bytes, deterministic per (o, blockIdx). */ + private fun q5_0Block(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(22) + le16(b, 0, 0x2C00 + rng.nextInt(0x80)) // small finite positive half + for (k in 2 until 22) b[k] = rng.nextInt(256).toByte() + return b + } + + private fun q5_1Block(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(24) + le16(b, 0, 0x2C00 + rng.nextInt(0x80)) // d + le16(b, 2, 0x2400 + rng.nextInt(0x80)) // m + for (k in 4 until 24) b[k] = rng.nextInt(256).toByte() + return b + } + + private fun q4_0Block(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(18) + le16(b, 0, 0x2C00 + rng.nextInt(0x80)) + for (k in 2 until 18) b[k] = rng.nextInt(256).toByte() + return b + } + + private fun q8_0Block(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(34) + le16(b, 0, 0x2C00 + rng.nextInt(0x80)) + for (k in 2 until 34) b[k] = rng.nextInt(256).toByte() + return b + } + + /** d @0, dMin @2 super-block scale header; ground truth only needs finite bytes. */ + private fun q4_kBlock(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(144) + le16(b, 0, 0x3400 + rng.nextInt(0x80)) + le16(b, 2, 0x2C00 + rng.nextInt(0x80)) + for (k in 4 until 144) b[k] = rng.nextInt(256).toByte() + return b + } + + private fun q5_kBlock(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(176) + le16(b, 0, 0x3400 + rng.nextInt(0x80)) + le16(b, 2, 0x2C00 + rng.nextInt(0x80)) + for (k in 4 until 176) b[k] = rng.nextInt(256).toByte() + return b + } + + /** Q6_K: `d` is the last two bytes of the 210-byte block. */ + private fun q6_kBlock(o: Int, blockIdx: Int, seed: Int): ByteArray { + val rng = Random(seed * 92821 + o * 977 + blockIdx) + val b = ByteArray(210) + for (k in 0 until 208) b[k] = rng.nextInt(256).toByte() + le16(b, 208, 0x3400 + rng.nextInt(0x80)) + return b + } + + /** + * Builds both the "canonical" (true row-major, output-major) and + * "kernel-native" (input-block-major) flat packings from the SAME set of + * per-(o,blockIdx) block contents. + */ + private fun buildPackings( + outputDim: Int, + blocksPerInputDim: Int, + bytesPerBlock: Int, + blockAt: (o: Int, blockIdx: Int) -> ByteArray, + ): Pair { + val canonical = ByteArray(outputDim * blocksPerInputDim * bytesPerBlock) + val kernelNative = ByteArray(outputDim * blocksPerInputDim * bytesPerBlock) + for (o in 0 until outputDim) { + for (bI in 0 until blocksPerInputDim) { + val block = blockAt(o, bI) + block.copyInto(canonical, (o * blocksPerInputDim + bI) * bytesPerBlock) + block.copyInto(kernelNative, (bI * outputDim + o) * bytesPerBlock) + } + } + return canonical to kernelNative + } + + private fun isAllZero(a: FloatArray): Boolean = a.all { it == 0f } + + private fun runGroundTruth( + name: String, + blockElems: Int, + bytesPerBlock: Int, + inputDim: Int, + outputDim: Int, + seed: Int, + build: (Shape, ByteArray) -> TensorData, + blockAt: (o: Int, blockIdx: Int) -> ByteArray, + ) { + val blocksPerInputDim = inputDim / blockElems + val (canonicalBytes, kernelNativeBytes) = buildPackings(outputDim, blocksPerInputDim, bytesPerBlock, blockAt) + + // Ground truth: block-sequential dequant of the CANONICAL bytes reconstructs + // W[o][i] row-major, independent of the matmul kernel under test. + val wStorage = build(Shape(outputDim, inputDim), canonicalBytes) as PackedBlockStorage + val wFlat = wStorage.toFloatArray() // row-major [outputDim][inputDim], length outputDim*inputDim + val rng = Random(seed + 1000) + val xf = FloatArray(inputDim) { rng.nextFloat() - 0.5f } + val yGroundTruth = FloatArray(outputDim) { o -> + var s = 0f + for (i in 0 until inputDim) s += xf[i] * wFlat[o * inputDim + i] + s + } + + 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 ctxPre = DirectCpuExecutionContext() + val wPre = ctxPre.fromData(build(Shape(inputDim, outputDim), kernelNativeBytes), FP32::class) + val xPre = ctxPre.fromFloatArray(Shape(1, inputDim), FP32::class, xf) + val yPre = ctxPre.ops.matmul(xPre, wPre).data.copyToFloatArray() + + // Tolerance relative to the output vector's overall scale (max |value|), + // not each element individually: outputs near a sign change have tiny + // |groundTruth[i]| where even small absolute FMA/summation-order noise + // balloons into a huge *per-element* relative error despite being + // consistent, small noise across the whole vector. Ground truth magnitudes + // range from O(1) (32-elem blocks) to O(1e3) (256-elem K-quant + // super-blocks summing 512 pseudo-random, largely-cancelling terms), so a + // fixed absolute tolerance would either reject correct large-scale results + // or accept near-zero-scale noise. A real dispatch bug (as originally + // observed) produces errors that are large a fraction of the vector's own + // scale (wrong sign, wrong magnitude on most/all outputs) — 1% of the + // vector's max magnitude comfortably separates "FMA reordering noise" from + // "reading the wrong bytes". + val scale = maxOf(yGroundTruth.maxOf { abs(it) }, 1e-6f) + fun matchesGroundTruth(y: FloatArray): Boolean = y.indices.all { i -> + abs(y[i] - yGroundTruth[i]) <= 1e-2f * scale + } + val classicMatchesGroundTruth = matchesGroundTruth(yClassic) + val preMatchesGroundTruth = matchesGroundTruth(yPre) + + println( + "[$name] blocksPerInputDim=$blocksPerInputDim " + + "classicAllZero=${isAllZero(yClassic)} classicMatchesGroundTruth=$classicMatchesGroundTruth " + + "preAllZero=${isAllZero(yPre)} preMatchesGroundTruth=$preMatchesGroundTruth " + + "groundTruth[0..3]=${yGroundTruth.take(4)} classic[0..3]=${yClassic.take(4)}", + ) + + // Regression contract (fixed): `ops.matmul(x, ops.transpose(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. + // Before the fix this failed for every format with blocksPerInputDim > 1 + // (all of them here except the dedicated single-block control case). + assertTrue(!isAllZero(yPre), "$name: pre-transposed native output must be non-zero") + assertTrue(preMatchesGroundTruth, "$name: pre-transposed native output must match ground truth") + assertTrue(!isAllZero(yClassic), "$name: classic (lazy-transpose) native output must be non-zero") + assertTrue( + classicMatchesGroundTruth, + "$name: classic (lazy-transpose) native output must match ground truth " + + "(SKaiNET-transformers#307 regression — transpose must physically reorder packed blocks)", + ) + } + + @Test + fun q5_0_ground_truth_multi_block_per_row() = runGroundTruth( + "Q5_0", blockElems = 32, bytesPerBlock = 22, inputDim = 256, outputDim = 16, seed = 201, + build = { s, b -> Q5_0BlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q5_0Block(o, bI, 201) }, + ) + + @Test + fun q5_1_ground_truth_multi_block_per_row() = runGroundTruth( + "Q5_1", blockElems = 32, bytesPerBlock = 24, inputDim = 256, outputDim = 16, seed = 202, + build = { s, b -> Q5_1BlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q5_1Block(o, bI, 202) }, + ) + + @Test + fun q4_0_ground_truth_multi_block_per_row() = runGroundTruth( + "Q4_0", blockElems = 32, bytesPerBlock = 18, inputDim = 256, outputDim = 16, seed = 203, + build = { s, b -> Q4_0BlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q4_0Block(o, bI, 203) }, + ) + + @Test + fun q8_0_ground_truth_multi_block_per_row() = runGroundTruth( + "Q8_0", blockElems = 32, bytesPerBlock = 34, inputDim = 256, outputDim = 16, seed = 204, + build = { s, b -> Q8_0BlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q8_0Block(o, bI, 204) }, + ) + + @Test + fun q4_k_ground_truth_multi_block_per_row() = runGroundTruth( + "Q4_K", blockElems = 256, bytesPerBlock = 144, inputDim = 512, outputDim = 12, seed = 206, + build = { s, b -> Q4_KBlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q4_kBlock(o, bI, 206) }, + ) + + @Test + fun q5_k_ground_truth_multi_block_per_row() = runGroundTruth( + "Q5_K", blockElems = 256, bytesPerBlock = 176, inputDim = 512, outputDim = 12, seed = 207, + build = { s, b -> Q5_KBlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q5_kBlock(o, bI, 207) }, + ) + + @Test + fun q6_k_ground_truth_multi_block_per_row() = runGroundTruth( + "Q6_K", blockElems = 256, bytesPerBlock = 210, inputDim = 512, outputDim = 12, seed = 208, + build = { s, b -> Q6_KBlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q6_kBlock(o, bI, 208) }, + ) + + // Single-block-per-row control: blocksPerInputDim == 1 collapses the + // "input-block-major" and "output-major" orderings into the SAME physical + // layout (block index reduces to `o` either way), so this is expected to + // pass even if the general (multi-block-per-row) case is broken above. + @Test + fun q5_0_ground_truth_single_block_per_row_control() = runGroundTruth( + "Q5_0-singleblock", blockElems = 32, bytesPerBlock = 22, inputDim = 32, outputDim = 16, seed = 205, + build = { s, b -> Q5_0BlockTensorData(s, b) as TensorData }, + blockAt = { o, bI -> q5_0Block(o, bI, 205) }, + ) +}