From 6050e9964600fd65a01f65fa6cc0127328b5308e Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Fri, 14 Aug 2026 20:36:39 +0200 Subject: [PATCH 1/2] fix: Q6_K/Q5_1/Q5_0 matmul silently falls through for MemorySegment-backed activations chooseQuantizedMatmulHeap (the fallback dispatcher DefaultCpuOpsJvm. chooseQuantizedMatmul intentionally defers to for Q6_K/Q5_1/Q5_0) required the activation tensor to be exactly FloatArrayTensorData, silently returning null otherwise. Real forward-pass activations produced with DirectCpuExecutionContext(tensorDataFactory = MemorySegmentTensorDataFactory()) - the configuration every production caller uses - are MemorySegmentTensorData, not FloatArrayTensorData, so they never matched. Dispatch fell through the unguarded matmulGeneric fallback, which has no packed-quant handling and threw ClassCastException: Byte cannot be cast to Float reading raw packed bytes. Use the universal TensorData.copyToFloatArray() instead of a type-gated cast; FloatArrayTensorData already overrides it with a cheap buffer.copyOf(), so the fast path is unaffected. Also flattens leading batch/sequence dimensions in DefaultCpuOpsJvm.chooseQuantizedMatmul before its own rank-2 fast-path check, since linearProject's `[..., in]` inputs are a real (if not the triggering) case that fast path can't otherwise reach. Fixes #991. --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 26 +++++++- .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 21 +++++- .../tensor/ops/QuantizedMemSegMatmulTest.kt | 64 +++++++++++++++++++ 3 files changed, 107 insertions(+), 4 deletions(-) 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 fc9c123df..0231e5297 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 @@ -528,9 +528,31 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory * whole set on non-JVM) resolve here. */ protected fun chooseQuantizedMatmulHeap(a: Tensor, b: Tensor): Tensor? { - if (a.dtype != FP32::class || a.shape.rank != 2 || b.shape.rank != 2) return null + if (a.dtype != FP32::class || b.shape.rank != 2 || a.shape.rank < 2) return null + if (a.shape.rank == 2) return chooseQuantizedMatmulHeap2D(a, b) + + // Attention linear projections legitimately pass `[..., in]` (see linearProject's kdoc) — + // flatten the leading batch/sequence dims into one so the specialized quant kernels below + // (which only understand `[batch, in]`) still get used, instead of silently falling + // through to matmulGeneric, which has no packed-quant handling at all (see SKaiNET#991). + val leading = a.shape.dimensions.copyOf(a.shape.rank - 1) + val flatBatch = leading.fold(1) { acc, d -> acc * d } + val inputDim = a.shape.dimensions.last() + val a2d = reshape(a, Shape(intArrayOf(flatBatch, inputDim))) + val result2d = chooseQuantizedMatmulHeap2D(a2d, b) ?: return null + val outputDim = result2d.shape.dimensions.last() + return reshape(result2d, Shape(leading + outputDim)) + } + + private fun chooseQuantizedMatmulHeap2D(a: Tensor, b: Tensor): Tensor? { if (a.shape[1] != b.shape[0]) return null - val inputBuffer = (a.data as? FloatArrayTensorData<*>)?.buffer ?: return null + // Any TensorData exposes copyToFloatArray() (FloatArrayTensorData overrides it with a cheap + // buffer.copyOf(); everything else — e.g. MemorySegmentTensorData — uses the generic + // row-major default). The previous strict `as? FloatArrayTensorData` cast meant activations + // backed by anything else (e.g. MemorySegment-backed FP32, as SKaiNET-transformers' attention + // path produces) silently declined here, falling through to the unguarded matmulGeneric + // fallback for quant types routed to this function (Q6_K, Q5_1, Q5_0) — see SKaiNET#991. + val inputBuffer = a.data.copyToFloatArray() val batchSize = a.shape[0] val inputDim = a.shape[1] val outputDim = b.shape[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 ee43f3576..2ed4a2106 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 @@ -503,11 +503,28 @@ internal class DefaultCpuOpsJvm( private fun chooseQuantizedMatmul(a: Tensor, b: Tensor): Tensor? { // Input must be FP32 if (a.dtype != FP32::class) return null - if (a.shape.rank != 2) return null + if (b.shape.rank != 2) return null + if (a.shape.rank < 2) return null + if (a.shape.rank == 2) return chooseQuantizedMatmul2D(a, b) + + // Defensive: attention linear projections can in principle pass `[..., in]` (see + // linearProject's kdoc) — flatten any leading batch/sequence dims into one so the + // specialized quant kernels below (which only understand `[batch, in]`) still get used. + // Not exercised by the SKaiNET#991 repro itself (that input was already rank-2 — the + // actual bug there was chooseQuantizedMatmulHeap2D requiring FloatArrayTensorData, fixed + // separately in DefaultCpuOps.kt), but kept as a real robustness gap this closes too. + val leading = a.shape.dimensions.copyOf(a.shape.rank - 1) + val flatBatch = leading.fold(1) { acc, d -> acc * d } + val inputDim = a.shape.dimensions.last() + val a2d = reshape(a, Shape(intArrayOf(flatBatch, inputDim))) + val result2d = chooseQuantizedMatmul2D(a2d, b) ?: return null + val outputDim = result2d.shape.dimensions.last() + return reshape(result2d, Shape(leading + outputDim)) + } + private fun chooseQuantizedMatmul2D(a: Tensor, b: Tensor): Tensor? { val bData = b.data val bShape = b.shape - if (bShape.rank != 2) return null val batchSize = a.shape[0] val inputDim = a.shape[1] 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 8cd194fc4..40cf87756 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 @@ -17,6 +17,7 @@ import sk.ainet.lang.tensor.data.Q6_KTensorData import sk.ainet.lang.tensor.data.Q8MemorySegmentMarker import sk.ainet.lang.tensor.data.Q8MemorySegmentTensorData import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.data.MemorySegmentTensorDataFactory import sk.ainet.lang.types.FP32 import java.lang.foreign.Arena @@ -325,4 +326,67 @@ class QuantizedMemSegMatmulTest { assertEquals(Shape(batchSize, outputDim), result.shape) arena.close() } + + // ── Q6_K + MemorySegment-backed activation (SKaiNET#991) ────────────────── + + /** + * Regression test for SKaiNET#991. Real attention-layer activations produced + * by [DirectCpuExecutionContext] wired with [MemorySegmentTensorDataFactory] + * (the config every production caller uses — see KLlamaJava.loadGGUF in + * SKaiNET-transformers) are `MemorySegmentTensorData`, not + * `FloatArrayTensorData`. `chooseQuantizedMatmul` (this class) intentionally + * does not intercept Q6_K/Q5_1/Q5_0 — the comment above its `when(bData)` + * block says they're "handled in DefaultCpuOpsBase via the kernel registry" — + * but `DefaultCpuOpsBase.chooseQuantizedMatmulHeap` required + * `a.data as? FloatArrayTensorData<*>` and silently returned null for + * anything else, so those quant types fell all the way through to + * `matmulGeneric`, which has no packed-quant handling and threw + * `ClassCastException: class java.lang.Byte cannot be cast to class + * java.lang.Float` reading the raw packed bytes as if they were `Float`. + * + * Fixed by having `chooseQuantizedMatmulHeap` call the universal + * `TensorData.copyToFloatArray()` instead of requiring the + * `FloatArrayTensorData` subtype specifically. + */ + @Test + fun `Q6_K matmul with MemorySegment-backed FP32 activation does not throw and stays finite`() { + val inputDim = Q6_KTensorData.BLOCK_SIZE // exactly one block per row + val outputDim = 2 + val numBlocks = outputDim + + val weightBytes = ByteArray(numBlocks * Q6_KTensorData.BYTES_PER_BLOCK) { i -> (i and 0x3F).toByte() } + // Force a small, finite half-float scale per block (last 2 bytes of each + // 210-byte Q6_K block) so we don't synthesize a NaN/Inf scale — mirrors + // the safeguard in Q6KMatmulTest.randomQ6KBytes. 0x3C00 = 1.0f16. + for (block in 0 until numBlocks) { + val dOffset = block * Q6_KTensorData.BYTES_PER_BLOCK + 208 + weightBytes[dOffset] = 0x00.toByte() + weightBytes[dOffset + 1] = 0x3C.toByte() + } + @Suppress("UNCHECKED_CAST") + val weight: Tensor = VoidOpsTensor( + Q6_KBlockTensorData(Shape(numBlocks, inputDim), weightBytes) as TensorData, + FP32::class, + ) + + // MemorySegmentTensorDataFactory, not DenseTensorDataFactory — this is the + // one detail that reproduces the real bug. `fpTensor()` above (used by + // every other test in this file) goes through DenseTensorDataFactory and + // yields FloatArrayTensorData, which never exercised the broken path. + val memSegFactory = MemorySegmentTensorDataFactory() + val inputData = memSegFactory.fromFloatArray( + Shape(1, inputDim), FP32::class, FloatArray(inputDim) { (it + 1).toFloat() / inputDim }, + ) + val input: Tensor = VoidOpsTensor(inputData, FP32::class) + + val transposedWeight = ops.transpose(weight) + assertTrue(transposedWeight.data is Q6_KTensorData, "transpose must preserve Q6_K packed layout") + + val result = ops.matmul(input, transposedWeight) + + assertEquals(Shape(1, outputDim), result.shape) + for (v in result.data.copyToFloatArray()) { + assertTrue(v.isFinite(), "Q6_K matmul with MemorySegment-backed input produced a non-finite value: $v") + } + } } From b8c4158be40d2c70632897e69f4c09c7efa38ec8 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Fri, 14 Aug 2026 22:49:33 +0200 Subject: [PATCH 2/2] fix: quantized matmul dispatch drops rank-1 decode-step activations to a broken fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit chooseQuantizedMatmulHeap (DefaultCpuOps.kt) and chooseQuantizedMatmul (DefaultCpuOpsJvm.kt) both required a.shape.rank >= 2, so a rank-1 activation — the single-token hidden-state vector a real forward pass produces once incremental decode moves past the initial batched prefill — skipped the packed-quant kernel dispatch entirely and fell through to matmulGeneric's untyped per-element TensorData.get(). For a packed-quant (or PreTransposedWeight-wrapped, e.g. PreTransposedQ4_K) weight, that get() returns the raw packed byte rather than a dequantized Float, throwing "ClassCastException: Byte cannot be cast to Float" on every decode step past the first token — reproduced end-to-end against a real Llama-3.2-1B GGUF via SKaiNET-transformers' KLlamaJava. Both guards now admit rank-1 (`a.shape.rank < 1`, down from `< 2`); the existing leading-dim-flattening logic already generalizes correctly to it (empty leading dims, flatBatch=1) without further changes. Also hardens matmulGeneric itself as defense in depth: a lazily materialized copyToFloatArray() fallback if a TensorData's generic get() ever returns something other than the tensor's own dtype, instead of an unconditional unsafe `as Float` cast. Adds a regression test exercising a rank-1 FP32 activation against a Q4_K-packed weight end to end. Co-authored-by: Claude --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 47 +++++++++++++++++-- .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 12 +++-- .../tensor/ops/QuantizedMemSegMatmulTest.kt | 39 +++++++++++++++ 3 files changed, 89 insertions(+), 9 deletions(-) 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 0231e5297..4199f1dbf 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 @@ -528,13 +528,21 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory * whole set on non-JVM) resolve here. */ protected fun chooseQuantizedMatmulHeap(a: Tensor, b: Tensor): Tensor? { - if (a.dtype != FP32::class || b.shape.rank != 2 || a.shape.rank < 2) return null + if (a.dtype != FP32::class || b.shape.rank != 2 || a.shape.rank < 1) return null if (a.shape.rank == 2) return chooseQuantizedMatmulHeap2D(a, b) // Attention linear projections legitimately pass `[..., in]` (see linearProject's kdoc) — // flatten the leading batch/sequence dims into one so the specialized quant kernels below // (which only understand `[batch, in]`) still get used, instead of silently falling // through to matmulGeneric, which has no packed-quant handling at all (see SKaiNET#991). + // Also covers rank-1 activations (a single-token hidden-state vector during incremental + // decode, once the KV cache is warm and a matmul no longer runs against a batched + // prefill) — `leading` is then empty and `flatBatch` is 1, i.e. `[in]` promotes to + // `[1, in]` and the result squeezes back down to `[out]`, the same as `rank > 2` + // already did; previously this rank guard sent every post-prefill decode step straight + // to matmulGeneric's untyped per-element `TensorData.get()` path, which — for a + // packed-quant (or pre-transposed-marker-wrapped, e.g. PreTransposedQ4_K) weight — + // returns the raw packed byte, not a dequantized Float. val leading = a.shape.dimensions.copyOf(a.shape.rank - 1) val flatBatch = leading.fold(1) { acc, d -> acc * d } val inputDim = a.shape.dimensions.last() @@ -701,6 +709,35 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory return mapped } + // Safety net for TensorData implementations whose generic per-element get() doesn't + // return the tensor's own dtype — e.g. a packed-quant weight wrapped in a + // PreTransposedWeight marker (PreTransposedQ4_K/Q5_K/Q6_K/...), whose delegated get() + // surfaces the raw packed byte rather than a dequantized Float. The dispatchers above + // this fallback (chooseQuantizedMatmulHeap et al.) now route both prefill (rank > 2) and + // single-token decode (rank == 1) activations to the packed-quant kernels, so this should + // no longer be hit on that path — kept as defense in depth for any other TensorData + // implementation with the same gap, materializing the dequantized array lazily (once, + // only if actually needed) rather than up front for every matmulGeneric call. + var aFallback: FloatArray? = null + var bFallback: FloatArray? = null + + fun flatIndex(dims: IntArray, indices: IntArray): Int { + var offset = 0 + for (i in dims.indices) offset = offset * dims[i] + indices[i] + return offset + } + + fun floatAt(data: TensorData, dims: IntArray, indices: IntArray, isA: Boolean): Float { + val raw = data.get(*indices) + if (raw is Float) return raw + val fallback = if (isA) { + aFallback ?: data.copyToFloatArray().also { aFallback = it } + } else { + bFallback ?: data.copyToFloatArray().also { bFallback = it } + } + return fallback[flatIndex(dims, indices)] + } + val outData = dataFactory.init(outShape, a.dtype) { outIdx -> val (batchIdx, mIdx, nIdx) = when { aIs1D && bIs1D -> Triple(IntArray(0), -1, -1) @@ -730,22 +767,22 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory var k = 0 while (k < kA) { val av: Float = if (aIs1D) { - a.data.get(*intArrayOf(k)) as Float + floatAt(a.data, aDims, intArrayOf(k), isA = true) } else { val aIdx = IntArray(aRank) if (aBatchIdx.isNotEmpty()) aBatchIdx.copyInto(aIdx) aIdx[aRank - 2] = mIdx aIdx[aRank - 1] = k - a.data.get(*aIdx) as Float + floatAt(a.data, aDims, aIdx, isA = true) } val bv: Float = if (bIs1D) { - b.data.get(*intArrayOf(k)) as Float + floatAt(b.data, bDims, intArrayOf(k), isA = false) } else { val bIdx = IntArray(bRank) if (bBatchIdx.isNotEmpty()) bBatchIdx.copyInto(bIdx) bIdx[bRank - 2] = k bIdx[bRank - 1] = nIdx - b.data.get(*bIdx) as Float + floatAt(b.data, bDims, bIdx, isA = false) } acc += av * bv k++ 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 2ed4a2106..6fe862d1b 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 @@ -504,15 +504,19 @@ internal class DefaultCpuOpsJvm( // Input must be FP32 if (a.dtype != FP32::class) return null if (b.shape.rank != 2) return null - if (a.shape.rank < 2) return null + if (a.shape.rank < 1) return null if (a.shape.rank == 2) return chooseQuantizedMatmul2D(a, b) // Defensive: attention linear projections can in principle pass `[..., in]` (see // linearProject's kdoc) — flatten any leading batch/sequence dims into one so the // specialized quant kernels below (which only understand `[batch, in]`) still get used. - // Not exercised by the SKaiNET#991 repro itself (that input was already rank-2 — the - // actual bug there was chooseQuantizedMatmulHeap2D requiring FloatArrayTensorData, fixed - // separately in DefaultCpuOps.kt), but kept as a real robustness gap this closes too. + // Also covers rank-1 activations: a single-token hidden-state vector during incremental + // decode (post-prefill, once the KV cache is warm) — `leading` is then empty and + // `flatBatch` is 1, so `[in]` promotes to `[1, in]` and the result squeezes back down to + // `[out]`, same as `rank > 2` already did. Previously excluded by `a.shape.rank < 2`, + // which sent every such decode step straight to matmulGeneric's untyped per-element + // `TensorData.get()` path — for a pre-transposed packed-quant weight (e.g. + // PreTransposedQ4_K) that returns the raw packed byte, not a dequantized Float. val leading = a.shape.dimensions.copyOf(a.shape.rank - 1) val flatBatch = leading.fold(1) { acc, d -> acc * d } val inputDim = a.shape.dimensions.last() 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 40cf87756..aba082961 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 @@ -389,4 +389,43 @@ class QuantizedMemSegMatmulTest { assertTrue(v.isFinite(), "Q6_K matmul with MemorySegment-backed input produced a non-finite value: $v") } } + + // ── Q4_K + rank-1 (single-token decode) activation ───────────────────────── + + /** + * Regression test for the "Byte cannot be cast to Float" crash reported against a real + * EdgeTranslator run: `chooseQuantizedMatmul`/`chooseQuantizedMatmulHeap` both required + * `a.shape.rank >= 2`, returning null for a rank-1 activation and sending it straight to + * `matmulGeneric`'s untyped per-element `TensorData.get()`, which — for a packed-quant weight + * (in production, one wrapped by SKaiNET-transformers' `PreTransposedQ4_K` marker; a plain + * `Q4_KBlockTensorData` reproduces the same dispatch gap here) — returns the raw packed byte, + * not a dequantized Float. Real attention forward passes run FP32 batched (rank >= 2) during + * prefill but drop to a bare `[in]` hidden-state vector once the KV cache is warm and decoding + * proceeds one token at a time, which is exactly the shape this test exercises. + */ + @Test + fun `Q4_K matmul with rank-1 activation does not throw and stays finite`() { + val inputDim = Q4_KTensorData.BLOCK_SIZE // exactly one block per row + val outputDim = 2 + val numBlocks = outputDim + + val weightBytes = ByteArray(numBlocks * Q4_KTensorData.BYTES_PER_BLOCK) { i -> (i and 0x3F).toByte() } + @Suppress("UNCHECKED_CAST") + val weight: Tensor = VoidOpsTensor( + Q4_KBlockTensorData(Shape(numBlocks, inputDim), weightBytes) as TensorData, + FP32::class, + ) + val transposedWeight = ops.transpose(weight) + assertTrue(transposedWeight.data is Q4_KTensorData, "transpose must preserve Q4_K packed layout") + + // Rank 1, not rank 2 — a single-token hidden-state vector, not a `[1, in]` batch. + val input = fpTensor(Shape(inputDim), FloatArray(inputDim) { (it + 1).toFloat() / inputDim }) + + val result = ops.matmul(input, transposedWeight) + + assertEquals(Shape(outputDim), result.shape) + for (v in result.data.copyToFloatArray()) { + assertTrue(v.isFinite(), "Q4_K matmul with a rank-1 activation produced a non-finite value: $v") + } + } }