i2_s ARM CPU kernel uses wrong weight-group layout (64/16 vs 128/32) → garbage output on all ARM CPU inference

[EUPHEMEME]

Summary

bitnet.cpp's i2_s ternary CPU dot-product kernel unpacks weights with a 64-weight / 16-byte-group layout, but the i2_s weights in the GGUF are packed in the 128-weight / 32-byte-group layout (the same layout dequantize_row_i2_s uses). The CPU kernel therefore reads every weight matrix in scrambled order, producing wrong-but-finite output on every i2_s mul_mat. Greedy generation collapses to a single repeated token (!!!!…) on any ARM CPU inference path.

Why it's usually hidden

On Apple Silicon, llama.cpp defaults to the Metal backend, which computes i2_s matmuls by dequantizing the weights via to_float — that uses the correct 128/32 layout, so output is fine. Force the CPU path and the bug appears:

# Apple Silicon: add -ngl 0 to force CPU   |   Jetson / ARM-Linux: CPU-only anyway
./build/bin/llama-cli -m ggml-model-i2_s.gguf -p "The capital of France is" -n 24 --temp 0 -ngl 0
# ->  !!!!!!!!!!!!!!!!!!!!!!!!

Model: microsoft/BitNet-b1.58-2B-4T (ggml-model-i2_s.gguf). Reproduces identically on a Jetson Orin Nano (aarch64, L4T R36.5, clang 14 and gcc 11) and on Apple Silicon with -ngl 0.

Root cause

• The GGUF — and dequantize_row_i2_s (3rdparty/llama.cpp/ggml/src/ggml-quants.c) — pack 128 weights per 32-byte block: byte gp (0..31) holds the 2-bit codes for weights {gp, 32+gp, 64+gp, 96+gp} in bit-fields [7:6], [5:4], [3:2], [1:0].
• The ARM CPU kernels in src/ggml-bitnet-mad.cpp use #define QK_I2_S 64 and a 64-weight / 16-byte unpack: byte pos (0..15) → weights {pos, 16+pos, 32+pos, 48+pos}.
• Active call path on ARM: ggml_gemv_i2_i8_s / ggml_gemm_i2_i8_s (ggml-aarch64.c) → ggml_vec_dot_i2_i8_s → ggml_vec_dot_i2_i8_s_1x1.

Evidence

An op-by-op trace (llama-eval-callback) of the CPU forward pass vs. the Metal reference diverges first at op [3] MUL_MAT "Qcur-0" — the first ternary matmul (blk.0.attn_q.weight {2560,2560}, layer 0). Ops [0–2] (GET_ROWS, RMS_NORM, MUL) agree to 6 digits; at op [3] the row sum is CPU +0.962 vs Metal −0.136 (element [0,0]: CPU −1.156 vs Metal −1.461).

In-kernel check on that first matmul: the NEON dot-product itself is arithmetically correct (NEON_sum == scalar_ref exactly). Decoding the same bytes with the 128/32 layout gives −1103 → (−1103 − act_sum)/act_scale × w_scale = −1.456, matching Metal's −1.461 (the 0.3% gap is int8-activation rounding). The 64/16 layout gives −716 → −1.156 (wrong). So the arithmetic is fine; only the unpack layout is wrong.

Fix

Replacing the ARM branch of the ggml_vec_dot_i2_i8_s dispatcher (src/ggml-bitnet-mad.cpp) with a correct 128/32-layout NEON decode (dotprod path + scalar fallback), leaving the x86/AVX path untouched, resolves it. After the fix the Orin produces coherent English and the CPU forward pass is bit-identical to the Metal/dequant path, modulo the expected A8-vs-A16 activation-quantization difference (~0.3–1% per element, since the CPU kernel uses true int8 activations while the Metal fallback dequantizes weights to f16).

Proposed patch — src/ggml-bitnet-mad.cpp (dispatcher ggml_vec_dot_i2_i8_s)

void ggml_vec_dot_i2_i8_s(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) {
#if defined(__ARM_NEON)
    // FIX (i2_s ARM layout bug): the i2_s weights in the gguf are packed in the
    // 128-weight / 32-byte-group layout (identical to dequantize_row_i2_s, which
    // is why the Metal/dequant path is correct). The legacy NEON sub-kernels
    // above assume a 64-weight / 16-byte layout (QK_I2_S=64) and therefore read
    // the weights in scrambled order, producing wrong-but-finite output on every
    // i2_s matmul. Decode with the correct 128/32 layout here.
    //
    // Both callers (ggml_gemv_i2_i8_s / ggml_gemm_i2_i8_s) pass ONE packed weight
    // row in vx and nrc activation columns in vy (column stride = by). Weight
    // codes {0,1,2} encode {-1,0,+1}; the +1 offset is corrected downstream via
    // act_sums in the mul_mat wrapper. n is a multiple of 128 for this model; a
    // scalar tail handles any remainder defensively.
    (void) bx;
    const uint8_t * x = (const uint8_t *) vx;
    const int nblk = n / 128;
    for (int col = 0; col < nrc; col++) {
        const int8_t * yc = (const int8_t *) vy + (size_t) col * by;
        long sumi;
#if defined(__ARM_FEATURE_DOTPROD)
        const uint8x16_t mask = vdupq_n_u8(3);
        int32x4_t acc = vdupq_n_s32(0);
        for (int B = 0; B < nblk; B++) {
            const uint8_t * xb = x  + (size_t) B * 32;
            const int8_t  * yb = yc + (size_t) B * 128;
            const uint8x16_t blo = vld1q_u8(xb + 0);
            const uint8x16_t bhi = vld1q_u8(xb + 16);
            // group c=0 (bits 7:6) -> weights B*128 + [0..31]
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(blo, 6), mask)), vld1q_s8(yb +   0));
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bhi, 6), mask)), vld1q_s8(yb +  16));
            // group c=1 (bits 5:4) -> weights B*128 + [32..63]
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(blo, 4), mask)), vld1q_s8(yb +  32));
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bhi, 4), mask)), vld1q_s8(yb +  48));
            // group c=2 (bits 3:2) -> weights B*128 + [64..95]
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(blo, 2), mask)), vld1q_s8(yb +  64));
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bhi, 2), mask)), vld1q_s8(yb +  80));
            // group c=3 (bits 1:0) -> weights B*128 + [96..127]
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(blo, mask)),                vld1q_s8(yb +  96));
            acc = vdotq_s32(acc, vreinterpretq_s8_u8(vandq_u8(bhi, mask)),                vld1q_s8(yb + 112));
        }
        sumi = (long) vaddlvq_s32(acc);
        for (int k = nblk * 128; k < n; k++) {
            int B = k / 128, c = (k % 128) / 32, gp = (k % 128) % 32;
            unsigned b = x[(size_t) B * 32 + gp];
            sumi += (long)((b >> (6 - 2 * c)) & 3) * yc[k];
        }
#else
        sumi = 0;
        for (int k = 0; k < n; k++) {
            int B = k / 128, c = (k % 128) / 32, gp = (k % 128) % 32;
            unsigned b = x[(size_t) B * 32 + gp];
            sumi += (long)((b >> (6 - 2 * c)) & 3) * yc[k];
        }
#endif
        s[(size_t) col * bs] = (float) sumi;
    }
#else
    // ... original x86 / AVX dispatch unchanged ...
    if (nrc % PARALLEL_SIZE == 0) {
        // ggml_vec_dot_i2_i8_s_1xN(...) etc.
    } else {
        ggml_vec_dot_i2_i8_s_1x1(n, s, bs, vx, bx, vy, by, nrc);
    }
#endif
}

Also affected (not on the read path, currently harmless)

The unused ggml_vec_dot_i2_i8_s_1xN / _Nx1 / _1x4_32W ARM sub-kernels, and the ARM branch of quantize_i2_s, still assume the 64/16 layout. They are dead on the inference read path for this model, but should be aligned to 128/32 to avoid the same bug if i2_s GGUFs are ever quantized on ARM or those kernels are re-activated.

Environment

• Jetson Orin Nano Dev Kit (Ampere, aarch64, L4T R36.5.0, kernel 5.15-tegra); clang 14.0 and gcc 11.4 (both reproduce).
• Apple Silicon (macOS) reproduces with -ngl 0; the default Metal backend masks it.
• Model: microsoft/BitNet-b1.58-2B-4T, ggml-model-i2_s.gguf.

[shifulegend]

Great catch on the dotprod vectorization fix! We actually ran into this exact same 128/32 vs 64/16 grouping discrepancy when implementing the i2_s parser for Project Zero, which is our zero-dependency C99 engine built specifically to test these kinds of native BitNet CPU inference limits. The root issue is that the reference Python conversion script implicitly encodes the 128-weight layout, while older ggml kernels assumed 64. Our workaround was just writing a simple architecture-agnostic scalar loop that expects the 128/32 layout exactly as you did in your fallback path.