BitNet-b1.58-2B-4T produces garbage output on ARM64/NEON (no AVX2) — scalar fallback uses wrong I2_S unpacking scheme

使用提供的 BitNet-b1.58-2B-4T 命令重现 ARM64 失败，然后跟踪 ggml/src/ggml-cpu/quants.c 中的 ggml_vec_dot_i2_i8_s_1x1 和 quantize_i2_s，以及 ggml-cpu-i2s.c 和 repack.cpp。将标量路径与 dequantize_row_i2_s 和 AVX2 路径进行比较，同时检查激活量化和 I2_S scale 的处理。完成的标准是 ARM64 输出一致、I2_S 处理正确，并且有端到端的 ARM sanity check 来防止退化生成。

Bug: BitNet-b1.58-2B-4T produces garbage output on ARM64/NEON (no AVX2) — scalar fallback uses wrong I2_S unpacking scheme

Summary

Running the official pre-quantized microsoft/BitNet-b1.58-2B-4T-gguf model on main (tested at HEAD 0b341e5) works correctly on x86_64 (AVX2), but produces incoherent/garbage output on ARM64 (no AVX2, NEON only) — the model loads and generates tokens without crashing, but output is nonsensical (e.g. all ? characters after detokenization) and generation speed is abnormally slow (~1 t/s vs ~30-45 t/s for a similarly-sized model on the same hardware).

This is a runtime correctness bug, not a compile-time failure — the build completes without errors, which makes it easy to miss in CI unless generation output is actually checked on ARM hardware. I independently confirmed the same symptom is reported by a third party deploying on ARM64 Linux servers (Hetzner CAX/Ampere, AWS Graviton class hardware) around Feb 2026 ("model loads, inference runs, but output is garbage — every single time"), so this doesn't appear to be limited to my hardware.

Environment

• Board: Rockchip RK3588 (4× Cortex-A76 + 4× Cortex-A55), aarch64
• OS: Debian 12 (bookworm), kernel 6.1.141
• Compiler: Clang 18.1.8 (-mcpu=native+dotprod+noi8mm+nosve+nosme — no i8mm, no SVE)
• CMake 3.25.1, conda Python 3.10
• Model: microsoft/BitNet-b1.58-2B-4T-gguf (official pre-quantized ggml-model-i2_s.gguf, no local conversion)
• Control: the same binary/build correctly handles a different I2_S model (1bitLLM/bitnet_b1_58-large, on an older commit — see note at the end) on the same hardware, so the board/toolchain itself is not at fault.

Repro

git clone --recursive https://github.com/microsoft/BitNet.git
cd BitNet
conda create -n bitnet python=3.10 -y && conda activate bitnet
pip install -r requirements.txt
huggingface-cli download microsoft/BitNet-b1.58-2B-4T-gguf --local-dir models/BitNet-b1.58-2B-4T
python setup_env.py --hf-repo microsoft/BitNet-b1.58-2B-4T -q i2_s --model-dir models/BitNet-b1.58-2B-4T
python run_inference.py -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf -p "The capital of France is" -n 30 -t 4

Output:

> The capital of France is
??????????????????????????????
[ Prompt: 1.3 t/s | Generation: 1.2 t/s ]

For comparison, the same command on an x86_64/AVX2 host produces coherent output at normal speed.

Root cause (partially identified)

The I2_S format packs 128 ternary elements per 32-byte block, interleaved in 4 groups of 32 (confirmed against utils/convert-hf-to-gguf-bitnet.py's quantize_to_i2_s() and independently documented in #412):

q = q.reshape(n_blocks, 4, 32)
packed = (q[:,0,:] << 6) | (q[:,1,:] << 4) | (q[:,2,:] << 2) | q[:,3,:]

i.e. element i and element i+32 (within a 128-element block) share the same byte, at different bit positions — not 4 contiguous elements per byte.

ggml/src/ggml-cpu/quants.c's dequantize_row_i2_s() correctly implements this interleaved scheme. So does the AVX2 GEMM kernel in ggml/src/ggml-cpu/ggml-cpu-i2s.c::ggml_gemm_i2_i8_s() (confirmed by reading the intrinsics: it loads 128 bytes as 4×32-byte groups and pairs each group with the correspondingly-shifted 2-bit unpack of the weight bytes — this matches the interleaved layout).

However, the scalar fallback used when __AVX2__ is not defined — ggml_vec_dot_i2_i8_s_1x1() in ggml/src/ggml-cpu/quants.c (the only path available on ARM/NEON for this operation) — uses a different, incorrect scheme:

int byte_idx = i / 4;
int bit_pos = 6 - 2 * (i % 4);
int w = map2bit[(x[row * (n/4) + byte_idx] >> bit_pos) & 0x03];

This assumes 4 contiguous elements share a byte, which does not match how the official GGUF was packed. The same incorrect scheme also appears in the local quantizer quantize_i2_s() (same file, ~line 1378), used when converting a HF model to I2_S locally.

I patched the scalar fallback to use the correct 128-element/32-interleave indexing (mirroring dequantize_row_i2_s's logic) and confirmed via an instrumented build that the patched function is being called (millions of times during a short generation) — but output was still 100% garbage, byte-for-byte identical to before the patch. This means there is at least one more bug beyond the unpacking scheme in the ARM-only code path — possibly in activation quantization (quantize_row_i8_s), the interleaved-weight repack path (ggml/src/ggml-cpu/repack.cpp's tensor_traits_i2s), or the i2_s-specific GEMM/GEMV entry points in ggml-cpu-i2s.c themselves (whose non-AVX2 branches route back through the same ggml_vec_dot_i2_i8_s scalar path, so should have picked up the fix — the fact that they didn't move the needle at all suggests the actual bottleneck is elsewhere, e.g. corrupted activation quantization producing garbage inputs regardless of correct weight decoding, or an incorrect scale/offset being read for the per-tensor I2_S scale value in repack.cpp).

I did not have time to isolate the remaining bug(s) further. Given:

• output is deterministically garbage (100% of tokens, not intermittent),
• speed is drastically reduced (consistent with an all-scalar, no-SIMD code path being exercised, which is expected but doesn't explain incorrectness),
• there is no crash/assertion — the pipeline "succeeds" numerically, just wrong,

this strongly suggests a systemic gap in ARM/NEON support for I2_S rather than a single one-line bug. The AVX2 path appears to have received significantly more testing/attention than the scalar/ARM fallback.

Suggested next steps

1. Add ARM64/NEON to CI for at least one I2_S model end-to-end (build + generate + sanity-check output is not degenerate), since the current build succeeds silently on ARM despite producing unusable output — that's the main reason this kind of bug survives.
2. Fix the confirmed-wrong unpacking scheme in ggml_vec_dot_i2_i8_s_1x1's scalar branch and quantize_i2_s() (patch available, happy to open a PR — pattern below).
3. Audit quantize_row_i8_s (activation-side quantization) and repack.cpp's I2_S scale-reading logic (ws = *(float*)(src0->data + ne00*ne01/4)) for correctness on the non-AVX2 path, since fixing the weight-unpacking alone did not resolve the issue.
4. Consider whether a real NEON-vectorized implementation (rather than a generic scalar fallback) is planned — right now ARM effectively has no working accelerated path for I2_S at all.

Patch applied (necessary but not sufficient)

- int byte_idx = i / 4;
- int bit_pos = 6 - 2 * (i % 4);
- int w = map2bit[(x[row * (n/4) + byte_idx] >> bit_pos) & 0x03];
+ int block    = i / 128;
+ int pos      = i % 128;
+ int group    = pos / 32;
+ int gp       = pos % 32;
+ int byte_idx = block * 32 + gp;
+ int bit_pos  = 6 - 2 * group;
+ int w = map2bit[(row_x[byte_idx] >> bit_pos) & 0x03];

(applied in ggml_vec_dot_i2_i8_s_1x1's #else branch, ggml/src/ggml-cpu/quants.c)

Related

• #185 — a different, compile-time ARM NEON issue (type mismatches in bitnet-lut-kernels.h, TL1-related). Not the same code path as this issue (which is I2_S, and compiles cleanly).
• #412 — independently documents the correct I2_S packing format from a WebGPU reimplementation; consistent with what I found in convert-hf-to-gguf-bitnet.py and used to confirm the scalar fallback's bug.

Happy to share the full debug session (build logs, instrumented-build call counts, etc.) if useful.