feat(backends): JNI + Kotlin/Native faces of the fused lm_head kernel (#1150)

skainet-backends/skainet-backend-jni-cpu/native/skainet_jni.c

@@ -191,3 +191,34 @@ Java_sk_ainet_exec_kernel_jni_JniKernels_ternaryF32Gemv(
         skainet_ternary_f32_gemv(in, inputOffset, (const uint8_t*) w, weightByteOffset,
                                  inputDim, outputDim, out, outputOffset))
 }
+
+/*
+ * ternary_lmhead_stage1 (#1150): fused 4-plane BITNET_PLANES lm_head — the
+ * vendored NeoGPU kernel behind skainet_ternary_lmhead_stage1. The FP16 row
+ * scales live inside the weight buffer, so the shim derives the uint16_t*
+ * from the pinned weight array at rowScaleByteOffset (2-byte aligned by the
+ * Kotlin seam's contract). Pins three arrays like the bitnet shim above.
+ */
+JNIEXPORT void JNICALL
+Java_sk_ainet_exec_kernel_jni_JniKernels_ternaryLmheadStage1(
+    JNIEnv* env, jobject thiz,
+    jfloatArray input, jint inputOffset,
+    jbyteArray weight, jint planesByteOffset, jint planeStrideBytes, jint rowScaleByteOffset,
+    jint inputDim, jint outputDim,
+    jfloatArray output, jint outputOffset
+) {
+    (void) thiz;
+    jfloat* in = (*env)->GetPrimitiveArrayCritical(env, input, NULL);
+    jbyte* w = in ? (*env)->GetPrimitiveArrayCritical(env, weight, NULL) : NULL;
+    jfloat* out = w ? (*env)->GetPrimitiveArrayCritical(env, output, NULL) : NULL;
+    if (out) {
+        skainet_ternary_lmhead_stage1(
+            in, inputOffset,
+            (const uint8_t*) w, planesByteOffset, planeStrideBytes,
+            (const uint16_t*) ((const uint8_t*) w + rowScaleByteOffset), 0,
+            inputDim, outputDim, out, outputOffset);
+    }
+    if (out) (*env)->ReleasePrimitiveArrayCritical(env, output, out, 0);
+    if (w) (*env)->ReleasePrimitiveArrayCritical(env, weight, w, JNI_ABORT);
+    if (in) (*env)->ReleasePrimitiveArrayCritical(env, input, in, JNI_ABORT);
+}

skainet-backends/skainet-backend-jni-cpu/src/androidTest/kotlin/sk/ainet/exec/kernel/jni/JniKernelParityTest.kt

@@ -288,6 +288,49 @@ class JniKernelParityTest {
     @Test
     fun ternary_f32_parity_threaded_regime() = assertTernaryF32Parity(inputDim = 256, outputDim = 1024, seed = 22)
 
+    /**
+     * ternary_lmhead_stage1 (#1150): the fused 4-plane lm_head vs a local
+     * reimplementation of its contract, weights from a synthetic BITNET_PLANES
+     * buffer (8 planes of random codes + FP16 row scales pinned to 1.0).
+     */
+    @Test
+    fun ternary_lmhead_parity() {
+        val n = 16; val k = 256
+        val rng = Random(31)
+        val planeStride = n * k / 4
+        val scalesOffset = 8 * planeStride
+        val weight = ByteArray(scalesOffset + 2 * n)
+        for (i in 0 until scalesOffset) {
+            // pack four random codes {0,1,2} per byte
+            var b = 0
+            for (lane in 0 until 4) b = b or (rng.nextInt(3) shl (lane * 2))
+            weight[i] = b.toByte()
+        }
+        for (o in 0 until n) { // FP16 1.0 = 0x3C00 LE
+            weight[scalesOffset + o * 2] = 0x00
+            weight[scalesOffset + o * 2 + 1] = 0x3C
+        }
+        val input = FloatArray(k) { rng.nextFloat() - 0.5f }
+        val out = FloatArray(n)
+        JniKernels.ternaryLmheadStage1(input, 0, weight, 0, planeStride, scalesOffset, k, n, out, 0)
+        for (o in 0 until n) {
+            var want = 0.0
+            var w = 1.0
+            for (q in 0 until 4) {
+                val base = q * planeStride + o * (k / 4)
+                var dot = 0.0
+                for (i in 0 until k) {
+                    val code = ((weight[base + i / 4].toInt() and 0xFF) shr ((i % 4) * 2)) and 3
+                    dot += (code - 1) * input[i]
+                }
+                want += dot * w
+                w /= 3.0
+            }
+            val diff = abs(want.toFloat() - out[o])
+            assertTrue("[$o]: reference=$want jni=${out[o]} diff=$diff", diff <= 1e-3f)
+        }
+    }
+
     @Test
     fun smoke_roundtrip() {
         val input = floatArrayOf(1f, 2f, 3f, 4f)

skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniKernels.kt

@@ -119,6 +119,17 @@ public object JniKernels {
         output: FloatArray, outputOffset: Int,
     )
 
+    /**
+     * `ternary_lmhead_stage1` (#1150): fused 4-plane BITNET_PLANES lm_head — the vendored NeoGPU
+     * kernel. The FP16 row scales live inside [weight] at [rowScaleByteOffset] (2-byte aligned).
+     */
+    public external fun ternaryLmheadStage1(
+        input: FloatArray, inputOffset: Int,
+        weight: ByteArray, planesByteOffset: Int, planeStrideBytes: Int, rowScaleByteOffset: Int,
+        inputDim: Int, outputDim: Int,
+        output: FloatArray, outputOffset: Int,
+    )
+
     public external fun q40Matmul(
         input: FloatArray, inputOffset: Int,
         weight: ByteArray, weightByteOffset: Int,

skainet-backends/skainet-backend-jni-cpu/src/main/kotlin/sk/ainet/exec/kernel/jni/JniTernaryLmhead.kt

@@ -0,0 +1,40 @@
+package sk.ainet.exec.kernel.jni
+
+import sk.ainet.backend.api.kernel.TernaryLmheadNative
+import sk.ainet.backend.api.kernel.TernaryPlanesKernelPack
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+
+/**
+ * The vendored NeoGPU fused 4-plane lm_head as a [TernaryLmheadNative] (#1150) — the Android/JNI
+ * face of `skainet_ternary_lmhead_stage1`, sharing its C with the FFM and cinterop consumers.
+ *
+ * Like [JniTernaryF32Gemv], no capability split: the LUT kernel needs only baseline NEON, so the
+ * BASELINE `libskainet_jni.so` carries the full SIMD path on every arm64 device.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+public object JniTernaryLmhead : TernaryLmheadNative {
+
+    override val name: String get() = if (JniKernels.isLoaded) "neon" else "unloaded"
+
+    override fun lmheadStage1(
+        activation: FloatArray, activationOffset: Int,
+        weight: ByteArray, planesByteOffset: Int,
+        planeStrideBytes: Int, rowScaleByteOffset: Int,
+        inputDim: Int, outputDim: Int,
+        out: FloatArray, outOffset: Int,
+    ) {
+        JniKernels.ternaryLmheadStage1(
+            activation, activationOffset,
+            weight, planesByteOffset, planeStrideBytes, rowScaleByteOffset,
+            inputDim, outputDim,
+            out, outOffset,
+        )
+    }
+
+    /**
+     * Install this kernel into the dispatcher, or leave the decoding reference serving and say
+     * so — removing the AAR is a supported configuration, never a crash.
+     */
+    public fun install(warn: (String) -> Unit = { println("[skainet] $it") }): String =
+        TernaryPlanesKernelPack.install(if (JniKernels.isLoaded) this else null, warn = warn)
+}

skainet-backends/skainet-backend-native-cpu/src/nativeMain/kotlin/sk/ainet/exec/kernel/NativeKnTernaryLmhead.kt

@@ -0,0 +1,65 @@
+package sk.ainet.exec.kernel
+
+import kotlinx.cinterop.ExperimentalForeignApi
+import kotlinx.cinterop.addressOf
+import kotlinx.cinterop.reinterpret
+import kotlinx.cinterop.usePinned
+import sk.ainet.backend.api.kernel.TernaryLmheadNative
+import sk.ainet.backend.api.kernel.TernaryPlanesKernelPack
+import sk.ainet.kernels.cinterop.skainet_ternary_lmhead_stage1
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+
+/**
+ * Kotlin/Native face of the vendored NeoGPU fused 4-plane lm_head (#1150): calls
+ * `skainet_ternary_lmhead_stage1` through cinterop from the static archive — the linuxArm64
+ * (Pi-4/Cortex-A72) consumption path, where the archive's vendored file is pinned to
+ * `-march=armv8-a`. The FP16 row scales live inside the pinned weight array; the pointer is
+ * derived at `rowScaleByteOffset` (2-byte aligned by the seam's contract — unaligned uint16
+ * reads are legal on AArch64 anyway).
+ */
+@OptIn(ExperimentalForeignApi::class, ExperimentalMemoryApi::class)
+public object NativeKnTernaryLmhead : TernaryLmheadNative {
+
+    override val name: String get() = "cinterop"
+
+    override fun lmheadStage1(
+        activation: FloatArray, activationOffset: Int,
+        weight: ByteArray, planesByteOffset: Int,
+        planeStrideBytes: Int, rowScaleByteOffset: Int,
+        inputDim: Int, outputDim: Int,
+        out: FloatArray, outOffset: Int,
+    ) {
+        require(inputDim % 4 == 0) {
+            "NativeKnTernaryLmhead: inputDim must be a multiple of 4; got $inputDim"
+        }
+        if (outputDim == 0) return
+        if (inputDim == 0) {
+            // The C kernel writes rowScale * 0 per row; mirror without pinning empty arrays.
+            out.fill(0f, outOffset, outOffset + outputDim)
+            return
+        }
+        activation.usePinned { inPin ->
+            weight.usePinned { wPin ->
+                out.usePinned { outPin ->
+                    skainet_ternary_lmhead_stage1(
+                        inPin.addressOf(0),
+                        activationOffset,
+                        wPin.addressOf(0).reinterpret(),
+                        planesByteOffset,
+                        planeStrideBytes,
+                        wPin.addressOf(rowScaleByteOffset).reinterpret(),
+                        0,
+                        inputDim,
+                        outputDim,
+                        outPin.addressOf(0),
+                        outOffset,
+                    )
+                }
+            }
+        }
+    }
+
+    /** Install this kernel into the dispatcher — the archive is linked in, so always available. */
+    public fun install(warn: (String) -> Unit = {}): String =
+        TernaryPlanesKernelPack.install(this, warn = warn)
+}

skainet-backends/skainet-backend-native-cpu/src/nativeTest/kotlin/sk/ainet/exec/kernel/NativeKnTernaryLmheadParityTest.kt

@@ -0,0 +1,87 @@
+package sk.ainet.exec.kernel
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.TernaryCodec
+import kotlin.math.abs
+import kotlin.random.Random
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * Proves the Kotlin/Native cinterop path of the vendored fused lm_head (#1150):
+ * [NativeKnTernaryLmhead] against a local reimplementation of the 4-plane fused contract
+ * (`out[o] = rowScale[o] · Σ_q (1/3^q) · dot(in, plane_q row o)`). Weights come from the real
+ * codec ([TernaryCodec.encodeBitNetPlanes]) so the buffer geometry is the production one.
+ *
+ * Runs on the host archive AND under qemu-aarch64 (`-PcrossArm64=true`); the C kernel spawns its
+ * 4 pthreads at any output_dim, so every case exercises its threading through cinterop.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class NativeKnTernaryLmheadParityTest {
+
+    private fun reference(
+        weight: ByteArray, firstPlane: Int, n: Int, k: Int, input: FloatArray,
+    ): FloatArray {
+        val planeStride = n * k / 4
+        val scalesOffset = 8 * planeStride
+        val rowBytes = k / 4
+        return FloatArray(n) { o ->
+            val bits = (weight[scalesOffset + o * 2].toInt() and 0xFF) or
+                ((weight[scalesOffset + o * 2 + 1].toInt() and 0xFF) shl 8)
+            val scale = sk.ainet.lang.types.Fp16Codec.decode(bits)
+            var acc = 0.0
+            var w = 1.0
+            for (q in 0 until 4) {
+                val base = (firstPlane + q) * planeStride + o * rowBytes
+                var dot = 0.0
+                for (i in 0 until k) {
+                    val code = ((weight[base + i / 4].toInt() and 0xFF) shr ((i % 4) * 2)) and 3
+                    dot += (code - 1) * input[i]
+                }
+                acc += dot * w
+                w /= 3.0
+            }
+            (acc * scale).toFloat()
+        }
+    }
+
+    private fun assertParity(n: Int, k: Int, seed: Int) {
+        val rng = Random(seed)
+        val values = FloatArray(n * k) { (rng.nextFloat() - 0.5f) * 2f }
+        val weight = TernaryCodec.encodeBitNetPlanes(values, n, k)
+        val input = FloatArray(k) { rng.nextFloat() - 0.5f }
+        val planeStride = n * k / 4
+        val scalesOffset = 8 * planeStride
+
+        for (firstPlane in intArrayOf(0, 4)) {
+            val out = FloatArray(n)
+            NativeKnTernaryLmhead.lmheadStage1(
+                input, 0, weight,
+                planesByteOffset = firstPlane * planeStride,
+                planeStrideBytes = planeStride,
+                rowScaleByteOffset = scalesOffset,
+                inputDim = k, outputDim = n,
+                out = out, outOffset = 0,
+            )
+            val expected = reference(weight, firstPlane, n, k, input)
+            for (o in 0 until n) {
+                val diff = abs(expected[o] - out[o])
+                assertTrue(
+                    diff <= 1e-3f * maxOf(1f, abs(expected[o])),
+                    "planes $firstPlane..${firstPlane + 3} [$o]: reference=${expected[o]} cinterop=${out[o]}",
+                )
+            }
+        }
+    }
+
+    @Test fun small_head() = assertParity(n = 8, k = 64, seed = 1)
+
+    @Test fun bitnet_hidden_size() = assertParity(n = 96, k = 2560, seed = 2)
+
+    @Test fun zero_input_dim_zeros_output() {
+        val out = FloatArray(3) { 9f }
+        NativeKnTernaryLmhead.lmheadStage1(FloatArray(0), 0, ByteArray(8), 0, 0, 0, 0, 3, out, 0)
+        for (v in out) assertEquals(0f, v)
+    }
+}

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -1327,14 +1327,21 @@ public final class sk/ainet/lang/memory/TernaryCodec {
 	public static synthetic fun decode$default (Lsk/ainet/lang/memory/TernaryCodec;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BIIILjava/lang/Object;)[F
 	public final fun decodeBitNet ([BII)[F
 	public static synthetic fun decodeBitNet$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIILjava/lang/Object;)[F
+	public final fun decodeBitNetPlanes ([BIII)[F
+	public static synthetic fun decodeBitNetPlanes$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIIILjava/lang/Object;)[F
+	public final fun decodeBitNetPlanesRow ([BIII[FII)V
+	public static synthetic fun decodeBitNetPlanesRow$default (Lsk/ainet/lang/memory/TernaryCodec;[BIII[FIIILjava/lang/Object;)V
 	public final fun decodeTq1_0 ([BII)[F
 	public static synthetic fun decodeTq1_0$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIILjava/lang/Object;)[F
 	public final fun decodeTq2_0 ([BII)[F
 	public static synthetic fun decodeTq2_0$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIILjava/lang/Object;)[F
 	public final fun encode (Lsk/ainet/lang/tensor/storage/TensorEncoding;[F)[B
 	public final fun encodeBitNet ([F)[B
+	public final fun encodeBitNetPlanes ([FII)[B
 	public final fun encodeTq1_0 ([F)[B
 	public final fun encodeTq2_0 ([F)[B
+	public final fun planesRowScale ([BIIII)F
+	public static synthetic fun planesRowScale$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIIIILjava/lang/Object;)F
 }
 
 public final class sk/ainet/lang/memory/ViewsKt {
@@ -4690,6 +4697,38 @@ public final class sk/ainet/lang/tensor/data/BitNetB158TensorData$Companion {
 	public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/BitNetB158TensorData;
 }
 
+public final class sk/ainet/lang/tensor/data/BitNetPlanesTensorData : sk/ainet/lang/tensor/data/TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage {
+	public static final field Companion Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData$Companion;
+	public fun <init> (Lsk/ainet/lang/tensor/Shape;[B)V
+	public fun copyToFloatArray ()[F
+	public fun dequantizeBlock (I[FI)V
+	public fun get ([I)Ljava/lang/Float;
+	public synthetic fun get ([I)Ljava/lang/Object;
+	public fun getBlockCount ()I
+	public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder;
+	public fun getBlockSize ()I
+	public final fun getCols ()I
+	public fun getElementCount ()J
+	public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getPackedData ()[B
+	public fun getPackedView ()Lsk/ainet/lang/memory/TensorView;
+	public fun getPhysicalBytes ()J
+	public final fun getRows ()I
+	public fun getShape ()Lsk/ainet/lang/tensor/Shape;
+	public fun getView ()Lsk/ainet/lang/memory/TensorView;
+	public final fun rowScale (I)F
+	public fun set ([IF)V
+	public synthetic fun set ([ILjava/lang/Object;)V
+	public fun toFloatArray ()[F
+	public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage;
+	public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage;
+}
+
+public final class sk/ainet/lang/tensor/data/BitNetPlanesTensorData$Companion {
+	public final fun fromFloats (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData;
+	public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData;
+}
+
 public final class sk/ainet/lang/tensor/data/DenseFloatArrayTensorData : sk/ainet/lang/tensor/data/FloatArrayTensorData {
 	public fun <init> (Lsk/ainet/lang/tensor/Shape;[F)V
 	public fun copyToFloatArray ()[F
@@ -7417,6 +7456,20 @@ public final class sk/ainet/lang/tensor/storage/TensorEncoding$BITNET_B1_58 : sk
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class sk/ainet/lang/tensor/storage/TensorEncoding$BITNET_PLANES : sk/ainet/lang/tensor/storage/TensorEncoding {
+	public static final field INSTANCE Lsk/ainet/lang/tensor/storage/TensorEncoding$BITNET_PLANES;
+	public static final field PLANES I
+	public static final field ROW_SCALE_BYTES I
+	public final fun bufferBytes (II)I
+	public fun equals (Ljava/lang/Object;)Z
+	public fun getName ()Ljava/lang/String;
+	public fun hashCode ()I
+	public fun physicalBytes (J)Ljava/lang/Long;
+	public final fun planeStrideBytes (II)I
+	public final fun rowScalesByteOffset (II)I
+	public fun toString ()Ljava/lang/String;
+}
+
 public final class sk/ainet/lang/tensor/storage/TensorEncoding$DENSE_I8_ABSMAX : sk/ainet/lang/tensor/storage/TensorEncoding {
 	public static final field CODE_RANGE I
 	public static final field INSTANCE Lsk/ainet/lang/tensor/storage/TensorEncoding$DENSE_I8_ABSMAX;