From 6e012acd3187811bfd3cc0fef48a2970e2d17d22 Mon Sep 17 00:00:00 2001 From: michalharakal Date: Wed, 21 Oct 2026 19:04:29 +0200 Subject: [PATCH] fix(native-cpu): multi-row products through the native Q8_0/Q4_0/Q5_0/Q5_1 kernels no longer zero the earlier rows The FFM bindings of these four kernels staged a zero-filled output segment of outputOffset + outputDim floats and copied the whole reach back from index 0, so a product served row by row into one array (DefaultCpuOpsJvm) kept only its last row. Single-row decode steps were right; the quantize-at-load lane on an encoder (Laya on Q8_0 with the native provider) returned uniform probabilities. The bindings now copy back exactly their slice, as the K-quant bindings do. NativeQ8CanonicalWeightTransposedReproTest covers every packed format for 1, 7 and 50 rows and the SPI offset contract. --- CHANGELOG.md | 4 + .../exec/kernel/NativeQ4_0MatmulKernel.kt | 5 +- .../exec/kernel/NativeQ5_0MatmulKernel.kt | 5 +- .../exec/kernel/NativeQ5_1MatmulKernel.kt | 5 +- .../exec/kernel/NativeQ8_0MatmulKernel.kt | 5 +- ...iveQ8CanonicalWeightTransposedReproTest.kt | 152 ++++++++++++++++++ 6 files changed, 172 insertions(+), 4 deletions(-) create mode 100644 skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ8CanonicalWeightTransposedReproTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index afa6ced35..8372137d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- **Multi-row activations through the native Q8_0 / Q4_0 / Q5_0 / Q5_1 kernels returned zeros for every row but the last.** The FFM bindings of these four kernels staged a zero-filled output segment of `outputOffset + outputDim` floats and copied the whole reach back from index 0, so when `DefaultCpuOpsJvm` serves a `[rows, in] × packed [out, in]ᵀ` product row by row into one array, each call wiped the rows computed before it. A single-row product (a decode step) was right, which is why the GGUF decode paths never showed it; the token rows of an encoder on the quantize-at-load lane (Laya on Q8_0 with `skainet-backend-native-cpu` on the classpath) came back uniform. The bindings now copy back exactly the slice they produced, like the K-quant bindings always did. `NativeQ8CanonicalWeightTransposedReproTest` covers every packed format for 1, 7 and 50 rows against the dequant ground truth, with the native provider alone and with the default registry, and the SPI contract that a call at an output offset writes only its own row. + ## [0.62.0] - 2026-10-21 Headline: **JVM ↔ Kotlin/Wasm parity for compute and IO.** In 0.61 the SIMD128 sidecar served the diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4_0MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4_0MatmulKernel.kt index 718b49171..5ea6076ca 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4_0MatmulKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ4_0MatmulKernel.kt @@ -78,7 +78,10 @@ internal object NativeQ4_0MatmulKernel : Q4_0MatmulKernel { outputSeg, outputOffset, ) - MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, outputReachFloats) + // Only the slice this call produced goes back: the staging segment is zero below + // outputOffset, and copying the whole reach wiped the rows a caller had already + // written there when it loops over the rows of a multi-row activation. + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, outputOffset.toLong() * java.lang.Float.BYTES, output, outputOffset, outputDim) } } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt index bb6da5038..34ea8fc22 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_0MatmulKernel.kt @@ -80,7 +80,10 @@ internal object NativeQ5_0MatmulKernel : Q5_0MatmulKernel { outputSeg, outputOffset, ) - MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, outputReachFloats) + // Only the slice this call produced goes back: the staging segment is zero below + // outputOffset, and copying the whole reach wiped the rows a caller had already + // written there when it loops over the rows of a multi-row activation. + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, outputOffset.toLong() * java.lang.Float.BYTES, output, outputOffset, outputDim) } } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt index 69cbb897f..935cb2080 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ5_1MatmulKernel.kt @@ -80,7 +80,10 @@ internal object NativeQ5_1MatmulKernel : Q5_1MatmulKernel { outputSeg, outputOffset, ) - MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, outputReachFloats) + // Only the slice this call produced goes back: the staging segment is zero below + // outputOffset, and copying the whole reach wiped the rows a caller had already + // written there when it loops over the rows of a multi-row activation. + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, outputOffset.toLong() * java.lang.Float.BYTES, output, outputOffset, outputDim) } } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ8_0MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ8_0MatmulKernel.kt index 23ad5536b..f5be2ac1f 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ8_0MatmulKernel.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeQ8_0MatmulKernel.kt @@ -82,7 +82,10 @@ internal object NativeQ8_0MatmulKernel : Q8_0MatmulKernel { outputSeg, outputOffset, ) - MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, outputReachFloats) + // Only the slice this call produced goes back: the staging segment is zero below + // outputOffset, and copying the whole reach wiped the rows a caller had already + // written there when it loops over the rows of a multi-row activation. + MemorySegment.copy(outputSeg, ValueLayout.JAVA_FLOAT, outputOffset.toLong() * java.lang.Float.BYTES, output, outputOffset, outputDim) } } diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ8CanonicalWeightTransposedReproTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ8CanonicalWeightTransposedReproTest.kt new file mode 100644 index 000000000..4647ae50d --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeQ8CanonicalWeightTransposedReproTest.kt @@ -0,0 +1,152 @@ +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.assertEquals +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 + +/** + * Multi-row activations through the native (FFM) packed kernels. + * + * `DefaultCpuOpsJvm` serves a `[rows, in] × packed [out, in]ᵀ` product by calling the format's SPI + * kernel once per row with `outputOffset = row * out` into one shared output array. The Q8_0, + * Q4_0, Q5_0 and Q5_1 bindings staged a zero-filled segment of `outputOffset + out` floats and + * copied the whole reach back from index 0 — so every call wiped the rows computed before it and a + * multi-row product came back as zeros except for its last row. A single-row product (decode) was + * right, which is why the GGUF decode paths never saw it; an encoder's token rows (Laya, BERT) did. + * + * Two contracts, with the native provider alone and with the default registry: + * - `ops.matmulWeightTransposed(x, w)` on a canonical row-major packed weight matches the + * block-sequential dequant ground truth for 1, 7 and 50 rows, every format; + * - a direct SPI call at a non-zero `outputOffset` writes exactly its slice and nothing else. + */ +class NativeQ8CanonicalWeightTransposedReproTest { + + @BeforeTest + fun nativeAvailable() { + assertTrue(NativeKernelProvider.isAvailable(), "the native provider must be available for this test") + } + + @AfterTest + fun reset() = KernelRegistry.clearForTesting() + + private class Format(val name: String, val blockElems: Int, val bytesPerBlock: Int, val build: (Shape, ByteArray) -> TensorData, val block: (Random, ByteArray, Int) -> Unit) + + private fun le16(b: ByteArray, o: Int, h: Int) { + b[o] = (h and 0xFF).toByte(); b[o + 1] = ((h ushr 8) and 0xFF).toByte() + } + + @Suppress("UNCHECKED_CAST") + private val formats: List = listOf( + Format("q8_0", 32, 34, { s, b -> Q8_0BlockTensorData(s, b) as TensorData }) { rng, b, base -> le16(b, base, 0x2C00 + rng.nextInt(0x80)); for (k in 2 until 34) b[base + k] = rng.nextInt(256).toByte() }, + Format("q4_0", 32, 18, { s, b -> Q4_0BlockTensorData(s, b) as TensorData }) { rng, b, base -> le16(b, base, 0x2C00 + rng.nextInt(0x80)); for (k in 2 until 18) b[base + k] = rng.nextInt(256).toByte() }, + Format("q5_0", 32, 22, { s, b -> Q5_0BlockTensorData(s, b) as TensorData }) { rng, b, base -> le16(b, base, 0x2C00 + rng.nextInt(0x80)); for (k in 2 until 22) b[base + k] = rng.nextInt(256).toByte() }, + Format("q5_1", 32, 24, { s, b -> Q5_1BlockTensorData(s, b) as TensorData }) { rng, b, base -> le16(b, base, 0x2C00 + rng.nextInt(0x80)); le16(b, base + 2, 0x2400 + rng.nextInt(0x80)); for (k in 4 until 24) b[base + k] = rng.nextInt(256).toByte() }, + Format("q4_k", 256, 144, { s, b -> Q4_KBlockTensorData(s, b) as TensorData }) { rng, b, base -> le16(b, base, 0x3400 + rng.nextInt(0x80)); le16(b, base + 2, 0x2C00 + rng.nextInt(0x80)); for (k in 4 until 144) b[base + k] = rng.nextInt(256).toByte() }, + Format("q5_k", 256, 176, { s, b -> Q5_KBlockTensorData(s, b) as TensorData }) { rng, b, base -> le16(b, base, 0x3400 + rng.nextInt(0x80)); le16(b, base + 2, 0x2C00 + rng.nextInt(0x80)); for (k in 4 until 176) b[base + k] = rng.nextInt(256).toByte() }, + Format("q6_k", 256, 210, { s, b -> Q6_KBlockTensorData(s, b) as TensorData }) { rng, b, base -> for (k in 0 until 208) b[base + k] = rng.nextInt(256).toByte(); le16(b, base + 208, 0x3400 + rng.nextInt(0x80)) }, + ) + + private fun canonicalBytes(f: Format, outputDim: Int, inputDim: Int, seed: Int): ByteArray { + val blocks = inputDim / f.blockElems + val bytes = ByteArray(outputDim * blocks * f.bytesPerBlock) + for (o in 0 until outputDim) for (bI in 0 until blocks) f.block(Random(seed * 92821 + o * 977 + bI), bytes, (o * blocks + bI) * f.bytesPerBlock) + return bytes + } + + private fun checkWeightTransposed(label: String, f: Format, rows: Int, inputDim: Int, outputDim: Int, seed: Int) { + val bytes = canonicalBytes(f, outputDim, inputDim, seed) + val ctx = DirectCpuExecutionContext() + val wData = f.build(Shape(outputDim, inputDim), bytes) + val wFlat = (wData as PackedBlockStorage).toFloatArray() + val w = ctx.fromData(wData, FP32::class) + val rng = Random(seed + 7) + val xf = FloatArray(rows * inputDim) { rng.nextFloat() - 0.5f } + val x = ctx.fromFloatArray(Shape(rows, inputDim), FP32::class, xf) + val truth = FloatArray(rows * outputDim) + for (r in 0 until rows) for (o in 0 until outputDim) { + var s = 0f + for (i in 0 until inputDim) s += xf[r * inputDim + i] * wFlat[o * inputDim + i] + truth[r * outputDim + o] = s + } + val scale = maxOf(truth.maxOf { abs(it) }, 1e-6f) + repeat(2) { pass -> // the second pass reads the relayout cache + val y = ctx.ops.matmulWeightTransposed(x, w).data.copyToFloatArray() + assertEquals(truth.size, y.size, "$label ${f.name}: output size") + val worst = y.indices.maxOf { abs(y[it] - truth[it]) } + val zeroRows = (0 until rows).count { r -> (0 until outputDim).all { y[r * outputDim + it] == 0f } } + println("[$label ${f.name}] rows=$rows in=$inputDim out=$outputDim pass=$pass zeroRows=$zeroRows worst=${"%.3g".format(worst)} scale=${"%.3g".format(scale)}") + assertEquals(0, zeroRows, "$label ${f.name} rows=$rows pass=$pass: $zeroRows all-zero output rows") + assertTrue(worst <= 1e-2f * scale, "$label ${f.name} rows=$rows pass=$pass: worst |diff| $worst exceeds 1% of scale $scale") + } + } + + private fun allRows(label: String) { + for (f in formats) { + val inputDim = if (f.blockElems == 256) 512 else 768 + checkWeightTransposed(label, f, rows = 1, inputDim = inputDim, outputDim = 96, seed = 1) + checkWeightTransposed(label, f, rows = 7, inputDim = inputDim, outputDim = 96, seed = 2) + checkWeightTransposed(label, f, rows = 50, inputDim = inputDim, outputDim = 64, seed = 3) + } + } + + @Test + fun multi_row_products_match_ground_truth_with_the_native_provider_alone() { + KernelRegistry.clearForTesting() + KernelRegistry.register(NativeKernelProvider) + allRows("native-only") + } + + @Test + fun multi_row_products_match_ground_truth_with_the_default_registry() { + KernelRegistry.clearForTesting() + allRows("default") + } + + @Test + fun a_call_at_an_output_offset_writes_only_its_own_slice() { + val inputDim = 64 + val outputDim = 48 + val rng = Random(5) + val input = FloatArray(inputDim) { rng.nextFloat() - 0.5f } + val q8 = canonicalBytes(formats[0], outputDim, inputDim, 9) + val q4 = canonicalBytes(formats[1], outputDim, inputDim, 9) + val q50 = canonicalBytes(formats[2], outputDim, inputDim, 9) + val q51 = canonicalBytes(formats[3], outputDim, inputDim, 9) + val calls: List Unit>> = listOf( + "q8_0" to { out, off -> NativeQ8_0MatmulKernel.matmul(input, 0, q8, 0, inputDim, outputDim, out, off) }, + "q4_0" to { out, off -> NativeQ4_0MatmulKernel.matmul(input, 0, q4, 0, inputDim, outputDim, out, off) }, + "q5_0" to { out, off -> NativeQ5_0MatmulKernel.matmul(input, 0, q50, 0, inputDim, outputDim, out, off) }, + "q5_1" to { out, off -> NativeQ5_1MatmulKernel.matmul(input, 0, q51, 0, inputDim, outputDim, out, off) }, + ) + for ((name, call) in calls) { + val reference = FloatArray(outputDim) + call(reference, 0) + assertTrue(reference.any { it != 0f }, "$name: reference output is all zero") + val sentinel = 123.5f + val out = FloatArray(3 * outputDim) { sentinel } + call(out, outputDim) // the middle row + for (i in 0 until outputDim) { + assertEquals(sentinel, out[i], "$name: row before the call must be untouched at $i") + assertEquals(sentinel, out[2 * outputDim + i], "$name: row after the call must be untouched at $i") + assertEquals(reference[i], out[outputDim + i], "$name: the call's own row at $i") + } + } + } +}