Skip to content

Commit 0101041

Browse files
michalharakalclaude
michalharakal
and
claude
committed
feat(backend): first-class Q5_K packed matmul + ARM NEON kernels
Adds Q5_K as a packed in-kernel dequant-matmul format (previously Q5_K was only eagerly decoded to FP32 on load), mirroring the existing Q4_K plumbing, and hand-written ARM NEON paths for the native CPU kernels. Q5_K (256-elt / 176-byte super-block: d, dMin, 12 packed scales, 32-byte qh high-bit plane, 128-byte qs low nibbles; 5-bit code = lowNibble | (5th<<4)): - TensorEncoding.Q5_K; Q5_KTensorData / Q5_KBlockTensorData (5th-bit fold). - Q5KMatmulKernel SPI + matmulQ5K()/"Q5_K" in KernelProvider.supports(). - ScalarQ5_KMatmulKernel (commonMain/KN), PanamaVectorQ5_KMatmulKernel (JVM), native C skainet_q5k_matmul + NativeQ5KMatmulKernel (FFM); all registered. - DefaultCpuOps matmul dispatch + lazy-transpose branches. - StreamingGgufParametersLoader: Q5_K + Q6_K packed branches (a Q5_K_M GGUF now loads end-to-end instead of SKIP'ing most tensors). Tests: Q5_KBlockTensorData bit-exact vs DequantOps golden across blocks; native<->Panama<->scalar matmul parity; KernelSupportMatrixTest gate updated. ARM NEON (behind #if __ARM_NEON in skainet_simd.h; x86 keeps the scalar fallback, re-verified green): - fp32 (broadcast+vfmaq), q8_0 (widen int8->f32+vfmaq), q4k/q5k (nibble unpack + dual code/input accumulators; q5k folds the qh 5th bit via a runtime-count vshlq_u8). - CMake aarch64 branch: -march=armv8.2-a+fp16+dotprod (no +i8mm — A55 lacks it). Cross toolchain-aarch64.cmake + opt-in -PcrossArm64 gradle tasks; default x86 build unaffected. BOARD-VERIFY-PENDING: the NEON paths are aarch64-syntax-validated (clang --target=aarch64) but not executed (x86 host, no QEMU). Run the parity tests under qemu-aarch64 or on the SL2610 before relying on them. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 6d85369 commit 0101041

25 files changed

Lines changed: 1525 additions & 4 deletions

File tree

‎skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,12 @@ public interface KernelProvider {
7979
*/
8080
public fun matmulQ6K(): Q6KMatmulKernel? = null
8181

82+
/**
83+
* F32 × Q5_K matmul kernel exposed by this provider, or `null` if
84+
* this provider does not specialize Q5_K. Same fall-through pattern.
85+
*/
86+
public fun matmulQ5K(): Q5KMatmulKernel? = null
87+
8288
/**
8389
* F32 × Q5_1 matmul kernel exposed by this provider, or `null` if
8490
* this provider does not specialize Q5_1. Same fall-through pattern.
@@ -126,6 +132,7 @@ public interface KernelProvider {
126132
"Q8_0" -> matmulQ8_0() != null
127133
"Q4_0" -> matmulQ4_0() != null
128134
"Q6_K" -> matmulQ6K() != null
135+
"Q5_K" -> matmulQ5K() != null
129136
"Q5_1" -> matmulQ5_1() != null
130137
"Q5_0" -> matmulQ5_0() != null
131138
else -> false

‎skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Q5KMatmulKernel.kt‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package sk.ainet.backend.api.kernel
2+
3+
/**
4+
* F32 input × Q5_K-packed weights matrix-vector multiply, in canonical
5+
* ggml super-block layout.
6+
*
7+
* output[outputOffset + o] = Σ_j input[inputOffset + j] · dequant(weight[o, j])
8+
* for j ∈ [0, inputDim), o ∈ [0, outputDim)
9+
*
10+
* Block layout (256-element super-block, 176 bytes/block; see
11+
* [sk.ainet.lang.tensor.data.Q5_KTensorData] kdoc for the byte map):
12+
* - bytes 0..1 : `d` (super-block scale, FP16 LE)
13+
* - bytes 2..3 : `dMin` (super-block min-scale, FP16 LE)
14+
* - bytes 4..15 : 12 bytes of packed (6-bit scaleIdx, 6-bit minIdx) for
15+
* 8 sub-blocks via ggml's `get_scale_min_k4` mixing
16+
* (identical to Q4_K)
17+
* - bytes 16..47 : 32 bytes `qh` high-bit plane (the 5th bit of each code)
18+
* - bytes 48..175: 128 bytes of 4-bit low nibbles, *strided* in 4 groups of
19+
* 32 bytes (identical layout to Q4_K's `qs`)
20+
*
21+
* Per sub-block s ∈ 0..7:
22+
* `scale[s] = d * scaleIdx[s]`
23+
* `offset[s] = dMin * minIdx[s]`
24+
* per element: `code = lowNibble | (fifthBit << 4)` (0..31);
25+
* `dequant = code * scale[s] - offset[s]`
26+
*
27+
* The lazy-`dmin` accumulation trick (used by every well-tuned K-quant
28+
* kernel including ggml's reference) avoids subtracting `offset` per
29+
* element by tracking `Σ(input · code)` and `Σ(input)` per sub-block
30+
* and combining as `scale * codeSum − offset * inputSum` once.
31+
*
32+
* Implementations MUST NOT mutate `input` or `weight`. They MAY assume
33+
* the arrays do not alias each other or `output`. They MUST fully
34+
* write the `outputDim` floats starting at `output[outputOffset]`.
35+
*
36+
* Packed-weight row-major contract: `weight` holds blocks laid out
37+
* `(blockIdx * outputDim + o) * 176` for output row `o` and input
38+
* block index `blockIdx`. This matches `Q5_KBlockTensorData.packedData`.
39+
*
40+
* `inputDim` MUST be a multiple of 256 (the Q5_K block size).
41+
*/
42+
public interface Q5KMatmulKernel {
43+
/**
44+
* @param input FP32 input vector (single row).
45+
* @param inputOffset element offset into [input] where the row starts.
46+
* @param weight packed Q5_K bytes for the full `outputDim × inputDim` weight tensor.
47+
* @param weightByteOffset byte offset into [weight] where block (0, 0) starts.
48+
* @param inputDim contraction dimension (must be a multiple of 256).
49+
* @param outputDim number of output cells.
50+
* @param output FP32 output vector.
51+
* @param outputOffset element offset into [output] where the row starts.
52+
*/
53+
public fun matmul(
54+
input: FloatArray, inputOffset: Int,
55+
weight: ByteArray, weightByteOffset: Int,
56+
inputDim: Int, outputDim: Int,
57+
output: FloatArray, outputOffset: Int,
58+
)
59+
}

‎skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import sk.ainet.backend.api.kernel.KernelProvider
66
import sk.ainet.backend.api.kernel.Q4KMatmulKernel
77
import sk.ainet.backend.api.kernel.Q4_0MatmulKernel
88
import sk.ainet.backend.api.kernel.Q5_0MatmulKernel
9+
import sk.ainet.backend.api.kernel.Q5KMatmulKernel
910
import sk.ainet.backend.api.kernel.Q5_1MatmulKernel
1011
import sk.ainet.backend.api.kernel.Q6KMatmulKernel
1112
import sk.ainet.backend.api.kernel.Q8_0MatmulKernel
@@ -33,6 +34,7 @@ public object ScalarKernelProvider : KernelProvider {
3334
override fun matmulQ4_0(): Q4_0MatmulKernel = ScalarQ4_0MatmulKernel
3435
override fun matmulQ4K(): Q4KMatmulKernel = ScalarQ4_KMatmulKernel
3536
override fun matmulQ6K(): Q6KMatmulKernel = ScalarQ6_KMatmulKernel
37+
override fun matmulQ5K(): Q5KMatmulKernel = ScalarQ5_KMatmulKernel
3638
override fun matmulQ5_1(): Q5_1MatmulKernel = ScalarQ5_1MatmulKernel
3739
override fun matmulQ5_0(): Q5_0MatmulKernel = ScalarQ5_0MatmulKernel
3840
}

‎skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarQ5_KMatmulKernel.kt‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package sk.ainet.exec.kernel
2+
3+
import sk.ainet.backend.api.kernel.Q5KMatmulKernel
4+
5+
/**
6+
* Scalar reference [Q5KMatmulKernel] — commonMain, so Q5_K packed matmul works
7+
* on Kotlin/Native / JS / WASM, not only the JVM SIMD path.
8+
*
9+
* Q5_K super-block: 256 elements / 176 bytes, block-major `(blockIdx*outputDim+o)*176`:
10+
* `d`(f16) `dMin`(f16) 12 scale bytes (ggml `get_scale_min_k4` packing) 32 `qh`
11+
* high-bit bytes 128 `qs` low-nibble bytes. Each of the 8 sub-blocks (32 elts)
12+
* contributes `codeSum*scale - inputSum*offset`, with `scale = d*scaleIdx`,
13+
* `offset = dMin*minIdx`, and the 5-bit `code = lowNibble | (fifthBit << 4)`.
14+
* Math mirrors `DequantOps.dequantQ5KFromBytes` and the Q4_K kernel (only the
15+
* 5th-bit fold differs).
16+
*/
17+
public object ScalarQ5_KMatmulKernel : Q5KMatmulKernel {
18+
19+
private const val BLOCK_SIZE = 256
20+
private const val SUB_BLOCK = 32
21+
private const val BYTES_PER_BLOCK = 176
22+
private const val QH_OFFSET = 16
23+
private const val QS_OFFSET = 48
24+
25+
override fun matmul(
26+
input: FloatArray, inputOffset: Int,
27+
weight: ByteArray, weightByteOffset: Int,
28+
inputDim: Int, outputDim: Int,
29+
output: FloatArray, outputOffset: Int,
30+
) {
31+
require(inputDim % BLOCK_SIZE == 0) {
32+
"ScalarQ5_KMatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim"
33+
}
34+
if (outputDim == 0) return
35+
if (inputDim == 0) { for (o in 0 until outputDim) output[outputOffset + o] = 0f; return }
36+
val blocksPerInputDim = inputDim / BLOCK_SIZE
37+
val scaleIdx = IntArray(8)
38+
val minIdx = IntArray(8)
39+
40+
for (o in 0 until outputDim) {
41+
var acc = 0f
42+
for (blockIdx in 0 until blocksPerInputDim) {
43+
val blockBase = weightByteOffset + (blockIdx * outputDim + o) * BYTES_PER_BLOCK
44+
val d = decodeHalf(((weight[blockBase + 1].toInt() and 0xFF) shl 8) or (weight[blockBase].toInt() and 0xFF))
45+
val dMin = decodeHalf(((weight[blockBase + 3].toInt() and 0xFF) shl 8) or (weight[blockBase + 2].toInt() and 0xFF))
46+
47+
// ggml get_scale_min_k4 over the 12 scale bytes (identical to Q4_K).
48+
val sc = blockBase + 4
49+
for (sb in 0 until 4) {
50+
scaleIdx[sb] = weight[sc + sb].toInt() and 0x3F
51+
minIdx[sb] = weight[sc + sb + 4].toInt() and 0x3F
52+
}
53+
for (sb in 4 until 8) {
54+
val low4S = weight[sc + sb + 4].toInt() and 0x0F
55+
val high2S = (weight[sc + sb - 4].toInt() and 0xFF) ushr 6
56+
scaleIdx[sb] = low4S or (high2S shl 4)
57+
val low4M = (weight[sc + sb + 4].toInt() and 0xFF) ushr 4
58+
val high2M = (weight[sc + sb].toInt() and 0xFF) ushr 6
59+
minIdx[sb] = low4M or (high2M shl 4)
60+
}
61+
62+
val qhBase = blockBase + QH_OFFSET
63+
val qsBase = blockBase + QS_OFFSET
64+
val inBlockBase = inputOffset + blockIdx * BLOCK_SIZE
65+
for (groupJ in 0 until 4) {
66+
val qsRegion = qsBase + groupJ * 32
67+
// sub-block lo (low nibbles) then hi (high nibbles) of the same 32 bytes;
68+
// the 5th bit comes from qh[i], bit (2*groupJ + half).
69+
for (half in 0 until 2) {
70+
val sb = 2 * groupJ + half
71+
val bit = 2 * groupJ + half
72+
val inStart = inBlockBase + sb * SUB_BLOCK
73+
var codeSum = 0f
74+
var inputSum = 0f
75+
for (i in 0 until 32) {
76+
val b = weight[qsRegion + i].toInt() and 0xFF
77+
val low = if (half == 0) (b and 0x0F) else (b ushr 4)
78+
val fifth = ((weight[qhBase + i].toInt() and 0xFF) ushr bit) and 0x01
79+
val code = low or (fifth shl 4)
80+
val v = input[inStart + i]
81+
codeSum += v * code
82+
inputSum += v
83+
}
84+
acc += codeSum * (d * scaleIdx[sb]) - inputSum * (dMin * minIdx[sb])
85+
}
86+
}
87+
}
88+
output[outputOffset + o] = acc
89+
}
90+
}
91+
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ import sk.ainet.lang.tensor.data.Q4_KTensorData
1818
import sk.ainet.lang.tensor.data.Q4_KBlockTensorData
1919
import sk.ainet.lang.tensor.data.Q6_KTensorData
2020
import sk.ainet.lang.tensor.data.Q6_KBlockTensorData
21+
import sk.ainet.lang.tensor.data.Q5_KTensorData
22+
import sk.ainet.lang.tensor.data.Q5_KBlockTensorData
2123
import sk.ainet.lang.tensor.data.Q5_1TensorData
2224
import sk.ainet.lang.tensor.data.Q5_1BlockTensorData
2325
import sk.ainet.lang.tensor.data.Q5_0TensorData
@@ -333,6 +335,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
333335
private val q4_0Kernel by lazy { resolveProvider { it.matmulQ4_0() != null }?.matmulQ4_0() }
334336
private val q4kKernel by lazy { resolveProvider { it.matmulQ4K() != null }?.matmulQ4K() }
335337
private val q6kKernel by lazy { resolveProvider { it.matmulQ6K() != null }?.matmulQ6K() }
338+
private val q5kKernel by lazy { resolveProvider { it.matmulQ5K() != null }?.matmulQ5K() }
336339
private val q5_1Kernel by lazy { resolveProvider { it.matmulQ5_1() != null }?.matmulQ5_1() }
337340
private val q5_0Kernel by lazy { resolveProvider { it.matmulQ5_0() != null }?.matmulQ5_0() }
338341

@@ -367,6 +370,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
367370
is Q5_1TensorData -> q5_1Kernel?.let { k -> run(bd.packedData, k::matmul) }
368371
is Q5_0TensorData -> q5_0Kernel?.let { k -> run(bd.packedData, k::matmul) }
369372
is Q4_KTensorData -> q4kKernel?.let { k -> run(bd.packedData, k::matmul) }
373+
is Q5_KTensorData -> q5kKernel?.let { k -> run(bd.packedData, k::matmul) }
370374
is Q6_KTensorData -> q6kKernel?.let { k -> run(bd.packedData, k::matmul) }
371375
is Q8_0TensorData -> q8_0Kernel?.let { k -> run(bd.packedData, k::matmul) }
372376
is Q4_0TensorData -> q4_0Kernel?.let { k -> run(bd.packedData, k::matmul) }
@@ -598,6 +602,7 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
598602
@Suppress("UNCHECKED_CAST")
599603
when (val d = tensor.data) {
600604
is Q4_KTensorData -> return newTensor(Q4_KBlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)
605+
is Q5_KTensorData -> return newTensor(Q5_KBlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)
601606
is Q6_KTensorData -> return newTensor(Q6_KBlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)
602607
is Q5_1TensorData -> return newTensor(Q5_1BlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)
603608
is Q5_0TensorData -> return newTensor(Q5_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData<T, V>, tensor.dtype, tensor)

‎skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import sk.ainet.backend.api.kernel.Fp32MatmulKernel
55
import sk.ainet.backend.api.kernel.KernelProvider
66
import sk.ainet.backend.api.kernel.Q4KMatmulKernel
77
import sk.ainet.backend.api.kernel.Q4_0MatmulKernel
8+
import sk.ainet.backend.api.kernel.Q5KMatmulKernel
89
import sk.ainet.backend.api.kernel.Q5_0MatmulKernel
910
import sk.ainet.backend.api.kernel.Q5_1MatmulKernel
1011
import sk.ainet.backend.api.kernel.Q6KMatmulKernel
@@ -65,6 +66,9 @@ public object PanamaVectorKernelProvider : KernelProvider {
6566
override fun matmulQ6K(): Q6KMatmulKernel? =
6667
if (isAvailable()) PanamaVectorQ6_KMatmulKernel else null
6768

69+
override fun matmulQ5K(): Q5KMatmulKernel? =
70+
if (isAvailable()) PanamaVectorQ5_KMatmulKernel else null
71+
6872
private fun isVectorApiClassLoaded(): Boolean = runCatching {
6973
Class.forName("jdk.incubator.vector.FloatVector")
7074
Class.forName("jdk.incubator.vector.VectorSpecies")

0 commit comments

Comments
 (0)