Ternary f32 LUT gemv — vendored NeoGPU kernel: exact FP32×b1.58 matmul for baseline-NEON ARM (Pi-4/A72)

[michalharakal]

Goal

Vendor NeoGPU's ternary LUT NEON kernel (anjaustin/neogpu, MIT) into SKaiNET as an exact FP32-activation × BitNet-b1.58 matmul path for baseline-NEON ARM. Agreed with upstream in anjaustin/neogpu#1.

Why, versus our existing sdot-based bitnet_gemv (#1041):

• f32 activations → exact results. No I8-absmax requantization step ([S2.8] P6: F32→I8 absmax requant adapter (Forward scope, traced) + reference bitnet_gemv #1040), no ~1.5% quantization error. Weight decode is exact, so logits match the FP32 reference bit-for-bit (modulo summation order).
• Baseline NEON only (armv8-a+simd). No FEAT_DotProd required — this is the fast path for Cortex-A72 / Raspberry Pi 4-class targets where the current NEON kernel falls back to vmull_s8.
• Upstream measured (Pi-4, 4 threads, full-vocab lm_head): 96 ms / 6.78 GOPS LUT vs 51 ms / 3.20 GOPS int8 (with quant error) vs 175 ms / 0.93 GOPS F16 NEON.

The kernel is src/hs_ml_ternary_neon.c in NeoGPU: 338 lines, fully standalone (only stdint/string/pthread/arm_neon), NEON + exact scalar fallback branches, 256-entry × 4-float LUT (4 KB, L1-resident), vld1q_f32(LUT[w[b]]) + vfmaq_f32 inner loop.

Design decisions (verified against both codebases)

1. No new TensorEncoding. NeoGPU's sequential 2-bit payload is byte-identical to our BITNET_B1_58 payload: LUT ((b >> (lane*2)) & 3) − 1 ≡ TernaryCodec.encodeBitNet (code shl ((i%4)*2), code = value+1). Same element order, same bias.
2. Key on BITNET_B1_58 (in-band trailing LE FP32 scale, TernaryCodec.bitNetScale), not TernaryPacked (out-of-band scale, unrecoverable from a TensorView). The C kernel sees only the payload; Kotlin applies the scale to the output.
3. Dispatch priority is free. KernelDispatch.matmul checks the exact key first; registering under matmul(FP32 dense contiguous × FP32/BITNET_B1_58 blocked_row_major) wins over the requantize→bitnet_gemv path with zero dispatcher changes. Pack absent → today's behavior, unchanged.
4. Vendor verbatim (REUSE .license sidecar, SPDX-License-Identifier: MIT, SPDX-FileCopyrightText: 2024 NeoGPU 贡献者) + a thin house-style adapter skainet_ternary_f32.c. NeoGPU's pthread threading (4 threads over N≥512) stays — it is the measured configuration. LUT-init race closed via pthread_once in the adapter. MSVC has no pthreads → vendored file compiles only on non-MSVC; adapter carries a portable scalar fallback.
5. -march pin. Non-Apple aarch64 kernel builds use -march=armv8.2-a+fp16+dotprod; the vendored + adapter files get pinned to -march=armv8-a (set_source_files_properties) — the whole point is dotprod-less hardware.
6. Code-3 handling. Byte code 3 decodes to +2.0 in both NeoGPU's LUT and our decodeBitNet — parity-consistent. It is rejected at load-time repack (I2_S import), never in the kernel.
7. I2_S layout trap. Stock BitNet.cpp GGUF I2_S (type 36) uses a 64-weight group layout (code = (block[j%16] >> (6 − 2*(j/16))) & 3); NeoGPU's own converter emits the sequential layout. Import repacks group→sequential once at load (WeightForm / EncodingRequest.RequantizeTo), so both file flavors feed the same kernel.

Architecture proof points (this effort is also a validation exercise)

Beyond the kernel itself, this port is a deliberate stress test of SKaiNET's claim to be modular and scalable. Success means each of the following lands as a plug-in — encoding + codec + kernel pack — with zero changes to the DSL NN builders, the DAG/tape machinery, KernelDispatch, TensorOps, or any consumer code:

1. 新建 SIMD kernel → faster on real HW in eager mode. A third-party (vendored, MIT) NEON kernel reaches eager execution (DIRECT mode) purely via exact-key registration in KernelDispatch — no compilation step, no fusion pass, no runtime special-casing. Measured win expected on dotprod-less aarch64 (Pi-4/A72), where the current path falls back to vmull_s8.
2. 新建 quant type → memory saving. I2_S import kept packed as BITNET_B1_58 = 0.25 bytes/weight (+scale) instead of 4 bytes when FP32-widened (today's [S2.7] P6: ternary encodings TQ1_0, TQ2_0, BITNET_B1_58 with block spec + activation hint; reference decoder and parity fixtures generated from the descriptor #1033 behavior) — ~16× weight memory on every ternary tensor, gigabytes at BitNet-2B scale.
3. Novel research format → just another encoding. BITNET_PLANES (multi-plane trit residual, [ternary-f32] Phase 6: BITNET_PLANES encoding (multi-plane trit + FP16 row scale) + codec + lmhead-stage1 kernel pack #1150) shows a format invented outside SKaiNET slots in as encoding + codec + pack, and consumers (SKaiNET-transformers) select it declaratively via WeightForm — the network definition stays a stock dense layer.
4. Optionality / graceful degradation. Every pack is an optional artifact: absent → dispatch falls back to the existing requantize/reference path, bit-identical to today, no crash. install(null) warns and registers nothing (contract pinned by tests).
5. Third-party code, first-class licensing. Verbatim MIT vendoring with REUSE sidecars — external innovation is absorbed without forking or rewriting it into house code.

Phase 5's benchmark scenario + a loaded-weight-memory measurement are the evidence artifacts for 1 and 2; the parity/dispatch test suites are the evidence for 3 and 4.

Sub-issues

• [ternary-f32] Phase 1: vendor NeoGPU LUT kernel + adapter + CMake + JVM FFM downcall #1137 — Phase 1: vendor kernel + adapter + CMake (both native trees) + JVM FFM downcall + goldens
• [ternary-f32] Phase 2: TernaryF32GemvNative SPI + reference ViewKernel + TernaryF32KernelPack #1138 — Phase 2: TernaryF32GemvNative SPI + reference ViewKernel + TernaryF32KernelPack (pure commonMain)
• [ternary-f32] Phase 3: JNI (Android) + Kotlin/Native cinterop bridges + qemu parity #1139 — Phase 3: JNI (Android) + Kotlin/Native cinterop bridges + qemu parity lane
• [ternary-f32] Phase 4: GGUF I2_S (type 36) import with group→sequential repack, keep-packed BITNET_B1_58 #1140 — Phase 4: GGUF I2_S (type 36) import, group→sequential repack, keep-packed BitNetB158TensorData (independent of 1–3)
• [ternary-f32] Phase 5: benchmark scenario + docs #1141 — Phase 5: benchmark scenario + docs (incl. NeoGPU's "what didn't work" tuning notes)
• [ternary-f32] Phase 6: BITNET_PLANES encoding (multi-plane trit + FP16 row scale) + codec + lmhead-stage1 kernel pack #1150 — Phase 6: BITNET_PLANES encoding (multi-plane trit + FP16 row scale) + codec + lmhead_stage1 kernel pack — the lm_head fast path stays format-driven (plain matmul in DSL/DAG, dispatch by weight encoding), no new op/layer/fusion pass
• [ternary-f32] Final gate: licensing, REUSE & OSS-compliance close-out for the vendored NeoGPU/BitNet work #1166 — Final gate: licensing, REUSE & OSS-compliance close-out — repo-level (reuse lint in CI, byte-identical vendor verification, BitNet.cpp interpretation credit) and artifact-level (third-party notices in the published jar/klib/AAR that ship compiled NeoGPU code), plus the upstream courtesies agreed in Porting the ternary LUT kernel into SKaiNET (Kotlin Multiplatform ML runtime) — a few questions anjaustin/neogpu#1. Closes last.

关注-up candidates (file when reached): thread-count control / house threadpool for the proj kernel (deviates from verbatim vendoring), baseline--march variant of the full FFM kernels lib for Pi-4, TernaryPacked dispatch once views can carry an out-of-band scale. (The NeoGPU "used by" PR and upstream comment moved into #1166.)

Verification lanes

• :skainet-backends:skainet-backend-native-cpu:jvmTest — FFM + scalar/NEON goldens (incl. all 256 byte values)
• :skainet-backends:skainet-backend-native-cpu:linuxArm64Test -PcrossArm64=true — NEON under qemu-aarch64
• JniKernelParityTest on arm64 device/emulator
• ./gradlew generateKernelMatrix refresh
• End-to-end: microsoft/bitnet-b1.58-2B-4T GGUF (group) + NeoGPU-converted GGUF (sequential) vs FP32-widened baseline

Related

• Upstream agreement: Porting the ternary LUT kernel into SKaiNET (Kotlin Multiplatform ML runtime) — a few questions anjaustin/neogpu#1
• [S2.7] P6: ternary encodings TQ1_0, TQ2_0, BITNET_B1_58 with block spec + activation hint; reference decoder and parity fixtures generated from the descriptor #1033 (ternary tensors widen to FP32 — the I2_S import adds a keep-packed path for BITNET_B1_58 without expanding [S2.7] P6: ternary encodings TQ1_0, TQ2_0, BITNET_B1_58 with block spec + activation hint; reference decoder and parity fixtures generated from the descriptor #1033's TQ1/TQ2 scope)
• [S2.8] P6: F32→I8 absmax requant adapter (Forward scope, traced) + reference bitnet_gemv #1040 (I8 absmax requant adapter), [S2.9] P6: NEON bitnet_gemv kernel pack skainet-kernels-android-neon (JNI, optional artifact, reference fallback) #1041 (NEON bitnet_gemv) — the existing int8 path remains the portable fallback when the pack is not installed
• Application layer (SKaiNET-transformers): BitNet b1.58 support + NeoGPU fused 4-plane lm_head (tracking) SKaiNET-transformers#335 — BitNet architecture + RequantizeTo(BITNET_PLANES) weight-form wiring for output.weight + two-stage top-k decode