fix: Q6_K/Q5_1/Q5_0 matmul silently falls through for MemorySegment-backed activations

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt

@@ -528,9 +528,39 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
      * whole set on non-JVM) resolve here.
      */
     protected fun <T : DType, V> chooseQuantizedMatmulHeap(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V>? {
-        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 < 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()
+        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 <T : DType, V> chooseQuantizedMatmulHeap2D(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V>? {
         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]
@@ -679,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<T, V>, 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<T, V>(outShape, a.dtype) { outIdx ->
             val (batchIdx, mIdx, nIdx) = when {
                 aIs1D && bIs1D -> Triple(IntArray(0), -1, -1)
@@ -708,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++

skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt

@@ -503,11 +503,32 @@ internal class DefaultCpuOpsJvm(
     private fun <T : DType, V> chooseQuantizedMatmul(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V>? {
         // 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 < 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.
+        // 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()
+        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 <T : DType, V> chooseQuantizedMatmul2D(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V>? {
         val bData = b.data
         val bShape = b.shape
-        if (bShape.rank != 2) return null
 
         val batchSize = a.shape[0]
         val inputDim = a.shape[1]

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,106 @@ 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<FP32, Float> = VoidOpsTensor(
+            Q6_KBlockTensorData(Shape(numBlocks, inputDim), weightBytes) as TensorData<FP32, Float>,
+            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<FP32, Float>(
+            Shape(1, inputDim), FP32::class, FloatArray(inputDim) { (it + 1).toFloat() / inputDim },
+        )
+        val input: Tensor<FP32, Float> = 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")
+        }
+    }
+
+    // ── 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<FP32, Float> = VoidOpsTensor(
+            Q4_KBlockTensorData(Shape(numBlocks, inputDim), weightBytes) as TensorData<FP32, Float>,
+            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")
+        }
+    }
 }