BitNet b1.58 support + NeoGPU fused 4-plane lm_head (tracking)

[michalharakal]

Goal

Support BitNet b1.58 models (e.g. microsoft/bitnet-b1.58-2B-4T) end-to-end, including NeoGPU's fused 4-plane lm_head optimization — the application-layer half of SKaiNET-developers/SKaiNET#1136 (vendored NeoGPU ternary LUT kernel; upstream agreement in anjaustin/neogpu#1).

This is also an architecture validation (see the proof points in SKaiNET-developers/SKaiNET#1136): a whole new model class with an exotic lm_head format should require no changes to this repo's network DSL, fusion machinery, or runtimes — only a ModelFamily, a network def built from stock layers, a declarative WeightForm request, and sampling-level code.

Per SKaiNET physiology, the fast lm_head is format-driven, not op-driven: formats, codecs and kernel packs (incl. BITNET_PLANES, SKaiNET#1150) live in the engine, dispatch picks the fused kernel from the weight's encoding, and the DSL/DAG keep a plain dense/matmul. This repo owns only the model-level parts: choosing that format for output.weight at load, and the two-stage sampling.

Dependency

This repo consumes SKaiNET via the Maven BOM (sk.ainet:skainet-bom) — work here needs a SKaiNET release containing:

• the BITNET_B1_58 keep-packed GGUF loader + TernaryF32KernelPack (SKaiNET#1137–#1140)
• the BITNET_PLANES encoding + lmhead_stage1 kernel pack (SKaiNET#1150; symbol exported in SKaiNET#1137) — this repo has no native build of its own

Why the fused lm_head

Upstream NeoGPU measured on Pi-4/Cortex-A72: full-vocab lm_head in 96 ms / 6.78 GOPS (4-thread, fused 4-plane LUT) vs 51 ms / 3.20 GOPS int8 with ~1.5% quant error. The weight format is 4 ternary planes with weights {1, ⅓, ⅑, 1/27} + per-row FP16 scale — effectively "8 bits as four ternary digits" (planes 4–7 exist too, used only for candidate rescoring).

Sub-issues

• [bitnet] T1: BitNet architecture — ModelFamily, network def, packed-weight loader #336 — T1: BitNet architecture (ModelFamily, network def, packed-weight loader); baseline unfused correctness
• [bitnet] T2: lm_head via BITNET_PLANES weight format (WeightForm) + two-stage top-k decode #337 — T2: lm_head via BITNET_PLANES weight format (WeightForm request at load) + two-stage top-k decode — no new op/layer/fusion pass

Follow-up candidate: forwardBatched computes the full-vocab projection for every prefill position then narrows to the last row — last-row-only lm_head is an easy independent win.

Verification

• Fused vs unfused logits parity on a real BitNet GGUF
• Decode via skainet-cli producing coherent text at BitNet-2B4T dims (H=2560, F=6912, vocab≈128k)
• Reference target: ~1 tok/sec decode on Pi-4 class hardware (NeoGPU end-to-end measurement)

[michalharakal]

Status reconciliation (2026-08-31): T1 (#336) and T2 (#337) are implemented and shipped — PR #353's "BitNet fold-in" merged the whole module (llm-inference/bitnet/: NetworkDef/NetworkLoader/GGUFNameResolver/PackedGgufLoader/TwoStageDecode, ModelFamily.BITNET + detect, sampleFromCandidates in SamplingUtils, CLI branch, synthetic-GGUF + two-stage tests), on engine BOM 0.52.0 which carries everything the Dependency section asks for.

What actually remains of this tracking issue's scope is now split into narrower issues:

• [bitnet] Tied lm_head never engages the planes kernel on 2B4T — FP32-widened head, ~0.45 tok/s #357 — the headline promise on the flagship model: 2B4T ships no output.weight, so the tied lm_head never hits the planesLmHead lane and runs FP32-widened (~0.45 tok/s per the [bitnet] T1: BitNet architecture — ModelFamily, network def, packed-weight loader #336 measurement).
• [bitnet] Wire BitNetTwoStageDecode into the generation loop (pre-lm_head hidden-state hook) #358 — BitNetTwoStageDecode is unwired: no pre-lm_head hidden-state hook in OptimizedLLMRuntime, CLI decodes single-stage.
• [bitnet] Maturity gate: golden-token parity, model-gated smoke, entry-point parity + CLI knobs #360 — verification half of this issue's checklist: golden-token parity, model-gated smoke (smoke-models.json has no BitNet row; SKaiNET-transformers → engine 0.51: full engine-loader migration, MAPPED default, BitNet fold-in #353's "3/3" run was manual), entry-point parity, CLI knobs.
• [bitnet] Full #346 conformance: BitNetWeightLoader/RuntimeWeights naming + llm-runtime/kbitnet facade #359 — full Common structure and naming convention for model family modules #346 conformance (BitNet is the template's first greenfield instantiation).
• Engine side: SKaiNET#1240 — the ternary packs are not in 0.52's self-healing SPI, which is why this repo's CLI still carries explicit NativeTernaryF32GemvKernel.install() lines.

The "Verification" items here (fused-vs-unfused parity, coherent CLI decode, ~1 tok/s Pi-4 reference) map onto #357/#360. Suggest ticking the two sub-issue boxes and keeping this open until #357/#358/#360 close.

Full cross-repo tracker: BITNET-IMPLEMENTATION-TRACKER.md in the working notes.

[michalharakal]

Close-out — the tracking issue's scope is delivered end to end:

• T1 ([bitnet] T1: BitNet architecture — ModelFamily, network def, packed-weight loader #336) and T2 ([bitnet] T2: lm_head via BITNET_PLANES weight format (WeightForm) + two-stage top-k decode #337): shipped via PR SKaiNET-transformers → engine 0.51: full engine-loader migration, MAPPED default, BitNet fold-in #353, residuals in feat(bitnet): planes-requantize the tied lm_head — 2B4T engages the fused head (#357) #361/feat(bitnet): wire the two-stage lm_head decode into generation (#358) #362; the Common structure and naming convention for model family modules #346-conformant module shape landed in refactor(bitnet): full #346 conformance — BitNetWeightLoader, RuntimeWeights, kbitnet facade (#359) #363, chat entry point in feat(chat): BitNet chat support — correct template, eot_id stop, resolver entry (#360 E3) #371.
• The fused lm_head on the flagship model: the tied-2B4T head is planes-requantized ([bitnet] Tied lm_head never engages the planes kernel on 2B4T — FP32-widened head, ~0.45 tok/s #357/PR feat(bitnet): planes-requantize the tied lm_head — 2B4T engages the fused head (#357) #361) and the two-stage top-200 decode runs in the generation loop ([bitnet] Wire BitNetTwoStageDecode into the generation loop (pre-lm_head hidden-state hook) #358/PR feat(bitnet): wire the two-stage lm_head decode into generation (#358) #362) — 4.15 → 6.2 tok/s on macOS Apple Silicon.
• Architecture validation: held. The whole family arrived as a ModelFamily, a decoderTransformerNetwork call, a declarative WeightForm request and sampling-level code — no new op, layer, fusion pass, or dispatcher change; the engine side stayed encoding+codec+pack (SKaiNET#1136's proof points).
• Verification, beyond what this issue asked: greedy decode is token-for-token equal to bitnet.cpp and the HF BF16 reference on 2B4T (PR fix(bitnet): RoPE pairing is NEOX-style — greedy now matches bitnet.cpp and the BF16 reference exactly #365 — which also caught and fixed a real RoPE-pairing bug the coherent-looking output hid), gated by the committed parity fixture; model-gated smoke row in place.
• Not measured here: the ~1 tok/s Pi-4-class reference from NeoGPU's end-to-end numbers — all measurements are macOS Apple Silicon; an on-device run belongs to SKaiNET#1244 (the Android decode-sample leg).
• The prefill full-vocab-then-narrow follow-up remains a known, unfiled optimization.

Engine-side tracking (SKaiNET#1136) closes with its final gate (SKaiNET#1166) at the 0.52.0 release.