Repository navigation
i2_s ARM CPU kernel uses wrong weight-group layout (64/16 vs 128/32) → garbage output on all ARM CPU inference #585
Description
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: bytegp(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.cppuse#define QK_I2_S 64and a 64-weight / 16-byte unpack: bytepos(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.
Activity
shifulegend commentedon Jul 15, 2026
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.
shifulegend commentedon Aug 12, 2026
If you need a reliable fallback for ARM CPU inference while this gets upstreamed, you might want to look at https://github.com/shifulegend/project-zero. It's a zero-dependency C99 engine (no Python, just GCC+make) that runs BitNet b1.58 models natively. It handles the weight layouts independently and hits 36 tok/s on an older Xeon. We haven't benchmarked the Jetson Orin specifically yet, but it builds out of the box for ARM/aarch64 environments.