[bitnet] T2: lm_head via BITNET_PLANES weight format (WeightForm) + two-stage top-k decode

[michalharakal]

Part of the BitNet tracking issue (T2, depends on T1 and on SKaiNET's BITNET_PLANES encoding — SKaiNET-developers/SKaiNET#1136 Phase 6).

Approach — no new op, layer, or fusion pass

This deliberately stays inside standard SKaiNET physiology. The network def keeps its stock VoidDense("output", vocabSize, dim); the DAG keeps its ordinary matmul(x, transpose(output.weight)) node. The fused kernel is reached purely by weight format: the lm_head weight is loaded as TensorEncoding.BITNET_PLANES, and KernelDispatch's exact-key match routes the matmul to the skainet_ternary_lmhead_stage1-backed kernel. DIRECT mode and the compiled/optimized path both funnel through ops.matmul, so both get the fast path with zero per-mode hooks. (Earlier draft's LLMFusionPass rule / fused_lm_head GenericOperation / custom TernaryLmHead module: dropped.)

What remains here is genuinely application-level: choosing the format for the output weight, and the two-stage sampling algorithm.

Tasks

• Weight form wiring (in the T1 BitNet loader): request EncodingRequest.RequantizeTo(BITNET_PLANES) for output.weight — the plane encoding runs once at load through SKaiNET's WeightForm machinery, traced like any other conversion. Make it a loader option (default on for BitNet), so the exact single-plane path remains selectable.
• Two-stage decode (NeoGPU Stage 2): take top-200 candidates from the Stage-1 logits (planes 0–3), refine only those rows with planes 4–7 (weights 1/3⁴..1/3⁷) via the codec's row accessors. Needs a candidate-rescoring API next to sampleFromLogits in SamplingUtils / GenerateExtensions — the current contract is a full-vocab logits array. Keep the single-stage path as default until this lands.

Tests

• Logits parity on a real BitNet GGUF: BITNET_PLANES lm_head vs exact BITNET_B1_58 lm_head — Stage-1-only is an approximation (plane truncation ≈ 1.85% of row scale); assert top-k stability rather than bitwise equality, and that two-stage rescoring restores the exact top-1
• Decode-mode coverage: DIRECT and OPTIMIZED produce identical logits (both dispatch by format, no mode-specific code)
• Prefill unaffected (forwardBatched's full-vocab-then-narrow waste is a separate follow-up issue)

Notes

• Upstream driver reference: NeoGPU src/hs_ml_infer.c:1063–1175 (lmh_plane_w[8], LMH_CANDIDATES = 200)
• hs_ml_lmhead_stage1 always spawns 4 pthreads (no threshold) — appropriate at vocab scale

[michalharakal]

The packed load path is implemented and proven end to end (commit HEAD of feat/336-bitnet-architecture; full sweep 710 tests / 0 failures):

• BitNetPackedGgufLoader — a ternary I2_S GGUF loads through the engine's StreamingGgufParametersLoader: projections arrive packed (BitNetB158TensorData, 0.25 B/weight), output.weight as BITNET_PLANES (lossless for exactly-ternary sources — plane 0 captures the normalized codes with zero residual). i2sLayout knob (GROUP_128 default / GROUP_64 / SEQUENTIAL), planesLmHead switch, OUT_IN orientation.
• From there the physiology does the rest: stock bitnetNetwork, linearProject's matmul(x, transpose(W)) → engine Wᵀ marker → KernelDispatch by storage format — the vendored NeoGPU NEON kernels when NativeTernaryF32GemvKernel.install() / NativeTernaryLmheadKernel.install() ran, the decoding reference otherwise. Correctness never depends on the packs.
• One engine bug found and fixed on the way (SKaiNET#1181): isHeapPackedWeight was a hardcoded Q-format list, so transpose dense-copied packed ternary weights through a codes-returning get() → CCE. Now is PackedBlockStorage — format-driven like everything else.
• The money test (BitNetPackedGgufLoadTest): a synthetic BitNet I2_S GGUF (BitNet.cpp GROUP_128 payloads + trailer scales + llama.cpp KV metadata, written by the test itself) → packed load asserted by data type → OptimizedLLMRuntime logits match the FP32-widened load of the same file, with and without the native packs installed.

Remaining before close: the two-stage top-k decode (planes 4–7 rescoring via the codec's row accessors — a sampling-level API next to sampleFromLogits; single-stage stays default) and a real microsoft/bitnet-b1.58-2B-4T smoke once a model file is at hand. Note the branch depends on SKaiNET 0.41.0-SNAPSHOT from the mavenLocal lane until a release ships (and on #1181 being merged).

[michalharakal]

Two-stage decode landed (commit ca03a38; full sweep 713 tests / 0 failures) — the last implementable item of this issue without a real model file:

• BitNetTwoStageDecode — planes 0–3 stage-1 scan over the full vocab + exact rescoring of candidates. One improvement over NeoGPU's flat top-200: the stage-1 truncation of row r is bounded by rowScale(r)·(Σ_{p=4..7}3⁻ᵖ)·Σ|h|, computable per row — so topK() rescores every row whose upper bound reaches the k-th best lower bound, and the result equals the exact full-matmul top-k (pinned against a brute-force oracle). The maxCandidates cap (default 200, NeoGPU's LMH_CANDIDATES) restores the heuristic only when scores are pathologically packed.
• ScoredToken + sampleFromCandidates in llm-core, next to sampleFromLogits — the candidate-shaped sampler this issue asked for, same temperature semantics.
• Deliberately sampling-level, not dispatch-level: KernelDispatch's matmul over BITNET_PLANES stays the exact 8-plane product; approximation is an application decision, made visibly.
• Snapshot re-published from develop with #1181 merged; everything green against it.

Remaining before this issue closes: a real microsoft/bitnet-b1.58-2B-4T GGUF smoke (needs the model file — also confirms the GROUP_128 default), and optionally wiring BitNetTwoStageDecode into a runtime generation loop (needs a pre-lm_head hidden-state hook on OptimizedLLMRuntime — a separate, small design). The branch still needs a real SKaiNET release to come off 0.41.0-SNAPSHOT.

[michalharakal]

Re-scope note (2026-08-31): both tasks here are implemented — RequantizeTo(BITNET_PLANES) for output.weight behind the planesLmHead loader option (BitNetPackedGgufLoader.kt:93-101), and BitNetTwoStageDecode + sampleFromCandidates/ScoredToken beside sampleFromLogits (SamplingUtils.kt:154-190), with the exact-top-k bound pinned by test.

Two residuals split out:

• [bitnet] Wire BitNetTwoStageDecode into the generation loop (pre-lm_head hidden-state hook) #358 — the "separate, small design" flagged here: no generation loop calls the two-stage path (no pre-lm_head hidden-state hook in OptimizedLLMRuntime; CLI is single-stage).
• [bitnet] Tied lm_head never engages the planes kernel on 2B4T — FP32-widened head, ~0.45 tok/s #357 — on 2B4T the planes lane never engages at all: the checkpoint ships no output.weight and the tied head is FP32-widened.

Suggest closing this in favor of #357/#358, maintainer's call.

[michalharakal]

Closing — both tasks shipped (PR #353) and the flagged residual landed: the pre-lm_head hook + generateTwoStage wired the two-stage decode end to end (PR #362, #358), the tied-2B4T head engages the planes lane (PR #361, #357), and the decode-mode/logits-parity verification is pinned by the merged tests. Measured: 6.2 tok/s two-stage on 2B4T vs 4.15 pre-planes.