Skip to content

Commit 5e5cf0f

Browse files
michalharakalclaude
michalharakal
and
claude
committed
feat(q4_0): Panama SIMD kernel + reconcile MemSeg path to split layout
Adds `PanamaVectorQ4_0MatmulKernel` (JDK Vector API): per block, decode the FP16 scale, unpack the 16 code bytes into 32 sign-corrected floats in the canonical ggml split layout, then SIMD-FMA against the input window. Wired through `PanamaVectorKernelProvider.matmulQ4_0()` (priority 50), so `DefaultCpuOpsJvm`'s `q4_0MatmulKernel` now prefers it over the scalar floor on JDK 21+. Also fixes a latent layout bug: the existing JVM MemSegment Q4_0 path (`JvmQuantizedVectorKernels.dotQ4_0BlockMemSeg` and `Q4MemorySegmentTensorData` get/set/copyToFloatArray) used an *interleaved* nibble layout (code[2k]/[2k+1] from byte k), which does NOT match real GGUF Q4_0 weights (split layout: low nibbles → 0..15, high → 16..31, per `DequantOps.dequantQ4_0FromBytes`). This mismatch is the likely reason the Q4_0 MemSeg path was never exercised end-to-end. All three sites + the test encoder are reconciled to the split layout, so the MemSeg path now agrees with the heap `Q4_0BlockTensorData`, the scalar/Panama SPI kernels, and canonical ggml. Tests: PanamaVectorQ4_0MatmulKernelParityTest (scalar≈panama within FMA tolerance), QuantizedMemSegMatmulTest still green under split layout. apiCheck green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 6fdef91 commit 5e5cf0f

7 files changed

Lines changed: 264 additions & 16 deletions

File tree

‎skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,11 @@ public final class sk/ainet/exec/kernel/PanamaVectorQ4KMatmulKernel : sk/ainet/b
8181
public fun matmul ([FI[BIII[FI)V
8282
}
8383

84+
public final class sk/ainet/exec/kernel/PanamaVectorQ4_0MatmulKernel : sk/ainet/backend/api/kernel/Q4_0MatmulKernel {
85+
public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorQ4_0MatmulKernel;
86+
public fun matmul ([FI[BIII[FI)V
87+
}
88+
8489
public final class sk/ainet/exec/kernel/PanamaVectorQ8_0MatmulKernel : sk/ainet/backend/api/kernel/Q8_0MatmulKernel {
8590
public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorQ8_0MatmulKernel;
8691
public fun matmul ([FI[BIII[FI)V

‎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
@@ -4,6 +4,7 @@ import sk.ainet.backend.api.kernel.Bf16MatmulKernel
44
import sk.ainet.backend.api.kernel.Fp32MatmulKernel
55
import sk.ainet.backend.api.kernel.KernelProvider
66
import sk.ainet.backend.api.kernel.Q4KMatmulKernel
7+
import sk.ainet.backend.api.kernel.Q4_0MatmulKernel
78
import sk.ainet.backend.api.kernel.Q8_0MatmulKernel
89
import sk.ainet.exec.tensor.ops.JvmCpuBackendConfig
910

@@ -49,6 +50,9 @@ public object PanamaVectorKernelProvider : KernelProvider {
4950
override fun matmulQ8_0(): Q8_0MatmulKernel? =
5051
if (isAvailable()) PanamaVectorQ8_0MatmulKernel else null
5152

53+
override fun matmulQ4_0(): Q4_0MatmulKernel? =
54+
if (isAvailable()) PanamaVectorQ4_0MatmulKernel else null
55+
5256
private fun isVectorApiClassLoaded(): Boolean = runCatching {
5357
Class.forName("jdk.incubator.vector.FloatVector")
5458
Class.forName("jdk.incubator.vector.VectorSpecies")

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

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package sk.ainet.exec.kernel
2+
3+
import jdk.incubator.vector.FloatVector
4+
import jdk.incubator.vector.VectorOperators
5+
import jdk.incubator.vector.VectorSpecies
6+
import sk.ainet.backend.api.kernel.Q4_0MatmulKernel
7+
8+
/**
9+
* SIMD-vectorized FP32 × Q4_0 matmul on the JDK Vector API.
10+
*
11+
* Pipeline per 32-element block:
12+
* 1. Decode the 2-byte FP16 scale `d` once.
13+
* 2. Unpack the 16 code bytes into 32 sign-corrected floats (`nibble - 8`)
14+
* in a reusable scratch buffer, using the canonical ggml **split**
15+
* layout (low nibbles → elements 0..15, high nibbles → 16..31). The
16+
* nibble-pair-per-byte packing makes a fully-fused `ByteVector`
17+
* pipeline awkward, so this kernel keeps the scratch-then-FMA shape
18+
* (same approach as the legacy `JvmQuantizedVectorKernels` Q4_0 path).
19+
* 3. SIMD-FMA the scratch against the matching input window into a
20+
* lane-wise block accumulator, reduce across lanes, and fold `* d`
21+
* exactly once per block.
22+
*
23+
* Numerical equivalence with [ScalarQ4_0MatmulKernel] is within FMA +
24+
* reordered-reduction tolerance — the same bar the Q8_0 / Q4_K Panama
25+
* kernels use.
26+
*/
27+
public object PanamaVectorQ4_0MatmulKernel : Q4_0MatmulKernel {
28+
29+
private const val BLOCK_SIZE = 32
30+
private const val BYTES_PER_BLOCK = 18
31+
32+
private val floatSpecies: VectorSpecies<Float> = FloatVector.SPECIES_PREFERRED
33+
34+
override fun matmul(
35+
input: FloatArray, inputOffset: Int,
36+
weight: ByteArray, weightByteOffset: Int,
37+
inputDim: Int, outputDim: Int,
38+
output: FloatArray, outputOffset: Int,
39+
) {
40+
require(inputDim % BLOCK_SIZE == 0) {
41+
"PanamaVectorQ4_0MatmulKernel: inputDim must be a multiple of $BLOCK_SIZE; got $inputDim"
42+
}
43+
if (outputDim == 0) return
44+
if (inputDim == 0) {
45+
for (o in 0 until outputDim) output[outputOffset + o] = 0f
46+
return
47+
}
48+
val blocksPerInputDim = inputDim / BLOCK_SIZE
49+
val step = floatSpecies.length()
50+
val loopBound = floatSpecies.loopBound(BLOCK_SIZE)
51+
val codeBuf = FloatArray(BLOCK_SIZE)
52+
53+
for (o in 0 until outputDim) {
54+
var acc = 0f
55+
for (blockIdx in 0 until blocksPerInputDim) {
56+
val blockBase = weightByteOffset + (blockIdx * outputDim + o) * BYTES_PER_BLOCK
57+
// FP16 scale — two LE bytes.
58+
val dBits = (weight[blockBase].toInt() and 0xFF) or
59+
((weight[blockBase + 1].toInt() and 0xFF) shl 8)
60+
val d = halfToFloat(dBits)
61+
62+
// Split-layout unpack: low nibbles → 0..15, high → 16..31.
63+
val codesBase = blockBase + 2
64+
for (j in 0 until 16) {
65+
val b = weight[codesBase + j].toInt() and 0xFF
66+
codeBuf[j] = ((b and 0x0F) - 8).toFloat()
67+
codeBuf[16 + j] = ((b ushr 4) - 8).toFloat()
68+
}
69+
70+
val inputBase = inputOffset + blockIdx * BLOCK_SIZE
71+
var blockAccVec = FloatVector.zero(floatSpecies)
72+
var k = 0
73+
while (k < loopBound) {
74+
val inV = FloatVector.fromArray(floatSpecies, input, inputBase + k)
75+
val cV = FloatVector.fromArray(floatSpecies, codeBuf, k)
76+
blockAccVec = inV.fma(cV, blockAccVec)
77+
k += step
78+
}
79+
var blockAcc = blockAccVec.reduceLanes(VectorOperators.ADD)
80+
// Scalar tail (only if floatSpecies.length() doesn't divide 32 — rare).
81+
while (k < BLOCK_SIZE) {
82+
blockAcc += input[inputBase + k] * codeBuf[k]
83+
k++
84+
}
85+
acc += blockAcc * d
86+
}
87+
output[outputOffset + o] = acc
88+
}
89+
}
90+
91+
/** Same FP16 → FP32 conversion as [ScalarQ4_0MatmulKernel]. */
92+
private fun halfToFloat(hbits: Int): Float {
93+
val sign = (hbits and 0x8000) shl 16
94+
val exp = (hbits and 0x7C00) shr 10
95+
val mant = hbits and 0x03FF
96+
return when (exp) {
97+
0 -> {
98+
if (mant == 0) Float.fromBits(sign)
99+
else {
100+
var m = mant
101+
var e = -14
102+
while ((m and 0x400) == 0) {
103+
m = m shl 1
104+
e--
105+
}
106+
m = m and 0x3FF
107+
Float.fromBits(sign or ((e + 127) shl 23) or (m shl 13))
108+
}
109+
}
110+
31 -> Float.fromBits(sign or (0xFF shl 23) or (mant shl 13))
111+
else -> Float.fromBits(sign or ((exp - 15 + 127) shl 23) or (mant shl 13))
112+
}
113+
}
114+
}

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

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -549,13 +549,14 @@ internal object JvmQuantizedVectorKernels {
549549
// Read f16 scale
550550
val scale = halfToFloat(read2BytesLE(weightSeg, blockByteOffset))
551551

552-
// Unpack 16 packed bytes → 32 sign-corrected nibbles. Two
553-
// nibbles per byte load means half the byte traffic of the
554-
// straight scalar dot product.
552+
// Unpack 16 packed bytes → 32 sign-corrected nibbles in the
553+
// canonical ggml *split* layout: low nibbles decode elements
554+
// 0..15, high nibbles decode elements 16..31. (Matches
555+
// DequantOps.dequantQ4_0FromBytes and Q4_0BlockTensorData.)
555556
for (k in 0 until 16) {
556557
val b = weightSeg.get(JAVA_BYTE_LE, codesOffset + k.toLong()).toInt() and 0xFF
557-
codeBuf[2 * k] = (b and 0x0F).toFloat() - 8f
558-
codeBuf[2 * k + 1] = (b ushr 4).toFloat() - 8f
558+
codeBuf[k] = (b and 0x0F).toFloat() - 8f
559+
codeBuf[16 + k] = (b ushr 4).toFloat() - 8f
559560
}
560561

561562
// SIMD FMA dot product.

‎skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorQ4_0MatmulKernelParityTest.kt‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
package sk.ainet.exec.kernel
2+
3+
import kotlin.math.abs
4+
import kotlin.random.Random
5+
import kotlin.test.Test
6+
import kotlin.test.assertEquals
7+
import kotlin.test.assertFailsWith
8+
import kotlin.test.assertTrue
9+
10+
/**
11+
* Numerical parity tests for [PanamaVectorQ4_0MatmulKernel] against
12+
* [ScalarQ4_0MatmulKernel]. Both kernels apply the same FP16-scale
13+
* decode + `(nibble - 8)` dequant in the canonical ggml split layout;
14+
* differences come from FMA + reordered-reduction order only.
15+
*
16+
* Tolerance scales with the number of Q4_0 blocks processed: `1e-2 *
17+
* blocksPerInputDim`, clamped to a `1e-2` floor — mirrors the Q8_0
18+
* parity test convention.
19+
*/
20+
class PanamaVectorQ4_0MatmulKernelParityTest {
21+
22+
private val blockSize = 32
23+
private val bytesPerBlock = 18
24+
25+
/** Random Q4_0 packed bytes; scales clamped to a small positive FP16. */
26+
private fun randomQ4_0Bytes(blocksPerInputDim: Int, outputDim: Int, seed: Int): ByteArray {
27+
val rng = Random(seed)
28+
val numBlocks = blocksPerInputDim * outputDim
29+
val bytes = ByteArray(numBlocks * bytesPerBlock)
30+
rng.nextBytes(bytes)
31+
for (block in 0 until numBlocks) {
32+
val base = block * bytesPerBlock
33+
bytes[base + 0] = 0x00.toByte()
34+
bytes[base + 1] = 0x22.toByte() // FP16 0x2200 ≈ 7.6e-3
35+
}
36+
return bytes
37+
}
38+
39+
private fun assertParity(
40+
inputDim: Int,
41+
outputDim: Int,
42+
seed: Int,
43+
tolPerBlock: Float = 1e-2f,
44+
) {
45+
val blocksPerInputDim = inputDim / blockSize
46+
val rng = Random(seed)
47+
val input = FloatArray(inputDim) { rng.nextFloat() - 0.5f }
48+
val weight = randomQ4_0Bytes(blocksPerInputDim, outputDim, seed)
49+
val outScalar = FloatArray(outputDim)
50+
val outPanama = FloatArray(outputDim)
51+
52+
ScalarQ4_0MatmulKernel.matmul(input, 0, weight, 0, inputDim, outputDim, outScalar, 0)
53+
PanamaVectorQ4_0MatmulKernel.matmul(input, 0, weight, 0, inputDim, outputDim, outPanama, 0)
54+
55+
val tol = (tolPerBlock * blocksPerInputDim.coerceAtLeast(1)).coerceAtLeast(tolPerBlock)
56+
for (i in outScalar.indices) {
57+
val diff = abs(outScalar[i] - outPanama[i])
58+
assertTrue(
59+
diff <= tol,
60+
"mismatch at $i: scalar=${outScalar[i]} panama=${outPanama[i]} diff=$diff tol=$tol",
61+
)
62+
}
63+
}
64+
65+
@Test fun single_block_single_output_matches_scalar() =
66+
assertParity(inputDim = 32, outputDim = 1, seed = 1)
67+
68+
@Test fun single_block_multiple_outputs_matches_scalar() =
69+
assertParity(inputDim = 32, outputDim = 7, seed = 2)
70+
71+
@Test fun multiple_blocks_single_output_matches_scalar() =
72+
assertParity(inputDim = 256, outputDim = 1, seed = 3)
73+
74+
@Test fun llm_typical_attention_proj_matches_scalar() =
75+
assertParity(inputDim = 512, outputDim = 512, seed = 4)
76+
77+
@Test fun llm_typical_ffn_proj_matches_scalar() =
78+
assertParity(inputDim = 256, outputDim = 1024, seed = 5)
79+
80+
@Test fun rejects_non_block_aligned_input_dim() {
81+
assertFailsWith<IllegalArgumentException> {
82+
PanamaVectorQ4_0MatmulKernel.matmul(
83+
FloatArray(31), 0,
84+
ByteArray(bytesPerBlock), 0,
85+
31, 1,
86+
FloatArray(1), 0,
87+
)
88+
}
89+
}
90+
91+
@Test fun zero_input_dim_zeros_output() {
92+
val out = FloatArray(5) { 9f }
93+
PanamaVectorQ4_0MatmulKernel.matmul(
94+
FloatArray(0), 0,
95+
ByteArray(0), 0,
96+
0, 5,
97+
out, 0,
98+
)
99+
for (v in out) assertEquals(0f, v, "output should be zeroed for inputDim=0")
100+
}
101+
102+
@Test fun provider_returns_panama_q4_0_when_available() {
103+
val kernel = PanamaVectorKernelProvider.matmulQ4_0()
104+
if (PanamaVectorKernelProvider.isAvailable()) {
105+
assertTrue(
106+
kernel === PanamaVectorQ4_0MatmulKernel,
107+
"Provider must hand out the Panama Q4_0 kernel when available",
108+
)
109+
} else {
110+
assertEquals(null, kernel, "Provider must return null when Vector API unavailable")
111+
}
112+
}
113+
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ class QuantizedMemSegMatmulTest {
4747

4848
/**
4949
* Encode a single Q4_0 block: 32 float values -> 18 bytes (2 scale + 16 packed nibbles).
50+
* Uses the canonical ggml *split* layout: code[j] is the low nibble of
51+
* byte j, code[j+16] is the high nibble of byte j.
5052
*/
5153
private fun encodeQ4_0Block(values: FloatArray): ByteArray {
5254
require(values.size == 32)
@@ -62,8 +64,8 @@ class QuantizedMemSegMatmulTest {
6264
val out = ByteArray(18)
6365
out[0] = (scaleHalf and 0xFF).toByte()
6466
out[1] = ((scaleHalf shr 8) and 0xFF).toByte()
65-
for (i in 0 until 16) {
66-
out[2 + i] = ((codes[2 * i + 1] shl 4) or codes[2 * i]).toByte()
67+
for (j in 0 until 16) {
68+
out[2 + j] = ((codes[j + 16] shl 4) or codes[j]).toByte()
6769
}
6870
return out
6971
}

‎skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/tensor/data/Q4MemorySegmentTensorData.kt‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,11 @@ public interface Q4MemorySegmentMarker : MemorySegmentBackedData {
2727
*
2828
* Q4_0 block layout (18 bytes per 32 elements):
2929
* - 2 bytes: f16 scale (little-endian)
30-
* - 16 bytes: packed 4-bit codes (32 values, 2 per byte)
30+
* - 16 bytes: packed 4-bit codes (32 values) in the canonical ggml
31+
* *split* layout — low nibbles decode elements 0..15, high nibbles
32+
* decode elements 16..31.
3133
*
32-
* Dequantization: output[i] = (nibble[i] - 8) * scale
34+
* Dequantization: output[j] = (nibble[j] - 8) * scale
3335
*
3436
* The segment is arena-managed and 64-byte aligned for SIMD access.
3537
*/
@@ -52,20 +54,24 @@ public class Q4MemorySegmentTensorData(
5254
val flatIndex = calcFlatIndex(indices)
5355
val blockIdx = flatIndex / blockSize
5456
val elemIdx = flatIndex % blockSize
55-
val codesByteOffset = segmentByteOffset + blockIdx.toLong() * bytesPerBlock + 2 + (elemIdx / 2).toLong()
57+
// Split layout: elements 0..15 are low nibbles of bytes 0..15,
58+
// elements 16..31 are the high nibbles of the same bytes.
59+
val byteInBlock = if (elemIdx < 16) elemIdx else elemIdx - 16
60+
val codesByteOffset = segmentByteOffset + blockIdx.toLong() * bytesPerBlock + 2 + byteInBlock.toLong()
5661
val packedByte = segment.get(JAVA_BYTE, codesByteOffset).toInt() and 0xFF
57-
val code = if (elemIdx % 2 == 0) packedByte and 0x0F else packedByte ushr 4
62+
val code = if (elemIdx < 16) packedByte and 0x0F else packedByte ushr 4
5863
return code.toByte()
5964
}
6065

6166
override fun set(vararg indices: Int, value: Byte) {
6267
val flatIndex = calcFlatIndex(indices)
6368
val blockIdx = flatIndex / blockSize
6469
val elemIdx = flatIndex % blockSize
65-
val codesByteOffset = segmentByteOffset + blockIdx.toLong() * bytesPerBlock + 2 + (elemIdx / 2).toLong()
70+
val byteInBlock = if (elemIdx < 16) elemIdx else elemIdx - 16
71+
val codesByteOffset = segmentByteOffset + blockIdx.toLong() * bytesPerBlock + 2 + byteInBlock.toLong()
6672
val currentByte = segment.get(JAVA_BYTE, codesByteOffset).toInt() and 0xFF
6773
val newNibble = value.toInt() and 0x0F
68-
val updated = if (elemIdx % 2 == 0) {
74+
val updated = if (elemIdx < 16) {
6975
(currentByte and 0xF0) or newNibble
7076
} else {
7177
(currentByte and 0x0F) or (newNibble shl 4)
@@ -83,11 +89,14 @@ public class Q4MemorySegmentTensorData(
8389
val scale = halfToFloat((b1 shl 8) or b0)
8490
val elemsInBlock = minOf(blockSize, shape.volume - outIdx)
8591
for (i in 0 until elemsInBlock) {
86-
val codeOff = blockOff + 2 + (i / 2).toLong()
92+
// Split layout: i<16 → low nibble of byte i; i>=16 → high nibble of byte i-16.
93+
val byteInBlock = if (i < 16) i else i - 16
94+
val codeOff = blockOff + 2 + byteInBlock.toLong()
8795
val packedByte = segment.get(JAVA_BYTE, codeOff).toInt() and 0xFF
88-
val code = if (i % 2 == 0) packedByte and 0x0F else packedByte ushr 4
89-
result[outIdx++] = (code - 8).toFloat() * scale
96+
val code = if (i < 16) packedByte and 0x0F else packedByte ushr 4
97+
result[outIdx + i] = (code - 8).toFloat() * scale
9098
}
99+
outIdx += elemsInBlock
91100
}
92101
return result
93102
}

0 commit comments

Comments
 (0)