Skip to content

Fix i2_s ARM CPU kernel: correct 128/32 weight-group layout (fixes #585) - #586

Open
Alexander Hoffman (EUPHEMEME) wants to merge 1 commit into
microsoft:mainfrom
EUPHEMEME:fix/i2s-arm-weight-layout
Open

Alexander Hoffman (EUPHEMEME) wants to merge 1 commit into
microsoft:mainfrom
EUPHEMEME:fix/i2s-arm-weight-layout

Conversation

@EUPHEMEME

Copy link
Copy Markdown

Fixes #585.

What

The i2_s ternary CPU dot-product kernel unpacked weights with a 64-weight / 16-byte group layout, but the gguf packs them in the 128-weight / 32-byte layout (the one dequantize_row_i2_s uses). It read every weight matrix in scrambled order → wrong-but-finite output → garbage generation on all ARM CPU inference. On Apple Silicon the default Metal backend hides it (it dequantizes weights, so uses the correct layout); force CPU with -ngl 0 and it reproduces, exactly like Jetson / ARM-Linux.

Fix

Replace the ARM branch of the ggml_vec_dot_i2_i8_s dispatcher with a correct 128/32-layout NEON decode (dotprod path + scalar fallback). The x86/AVX path is unchanged.

Validation

  • First divergent op vs. the Metal reference was [3] MUL_MAT "Qcur-0" (first ternary matmul, layer 0). With the fix, the full CPU forward pass is bit-identical to the Metal/dequant path, modulo the expected int8-vs-f16 activation-quantization difference (~0.3–1% per element).
  • BitNet-b1.58-2B-4T on a Jetson Orin Nano (aarch64, clang 14 / gcc 11) goes from !!!! to coherent English.

See #585 for the full op-by-op trace and the in-kernel proof (the NEON arithmetic is correct; only the unpack layout was wrong).

Follow-up (not in this PR): the unused _1xN / _Nx1 / _1x4_32W ARM sub-kernels and the ARM branch of quantize_i2_s still assume the 64/16 layout — dead on the inference read path for this model, but worth aligning.

@EUPHEMEME

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree

@EUPHEMEME

Copy link
Copy Markdown
Author

Flagging that this is now an incomplete fix for main, and offering to extend it.

When this PR was opened, the ARM QK_I2_S assumption lived only in src/ggml-bitnet-mad.cpp. As of the currently pinned submodule (Eddie-Wang1120/llama.cpp @ 390c3077), the same

#elif defined(__ARM_NEON)
#include <arm_neon.h>
#define QK_I2_S 64

block now also appears in:

  • ggml/src/ggml-cpu/quants.c (lines 1294–1298)
  • ggml/src/ggml-cpu/ggml-cpu-i2s.c (lines 15–20)

Both carry the identical 128 / 64 / 128 pattern, so merging this PR as-is would fix one of three copies and leave ARM decoding against the wrong block size in the other two.

Happy to extend this PR to cover both submodule sites, or to split them into a separate PR against the submodule repo — whichever you'd prefer. Just let me know.

Separately, while testing this I found a live SIGSEGV on any i2_s model at n_ubatch >= 32 on builds that link BLAS (the default on macOS via Accelerate). It's independent of this PR — I'll open it as its own issue rather than widen the scope here.

@matheus-c-feitosa

Copy link
Copy Markdown

Verified on aarch64 against a reference: this PR fixes the output on its own base.

Setup: Neoverse-N1, clang 18.1.3, PR head a92dbc3 (on 01eb415, submodule 1f86f05), built the way setup_env.py does for i2_s, official ggml-model-i2_s.gguf, WikiText-2, 8 x 512 windows:

build greedy text PPL
01eb415 "!!!!!!!!!!!!" 40118134
01eb415 + this PR "Paris. Paris is a city..." 18.8515
HF transformers, same tokens "Paris. Paris is a city..." 18.7934

The result is identical, window by window, to 404980e
(June 2025, the last commit before 112f853).

One caveat: current main (0b341e5) no longer compiles this file. Since the submodule moved to 390c307, the I2_S kernels that run are in the submodule's ggml/src/ggml-cpu/quants.c, so the fix would have to go there (isHuangXin/llama.cpp#2 covers the layout part).

Analysis assisted by Claude Code; all runs were executed and all numbers checked by me.

The ARM CPU i2_s dot-product path decoded weights with a 64-weight / 16-byte
group layout (QK_I2_S=64), but the gguf packs i2_s weights in the 128-weight /
32-byte layout used by dequantize_row_i2_s. It therefore read every weight
matrix in scrambled order, producing wrong-but-finite output and garbage
generation on all ARM CPU inference (Jetson / ARM-Linux, and Apple Silicon with
-ngl 0; the default Metal path masks it).

Replace the ARM branch of the ggml_vec_dot_i2_i8_s dispatcher with a correct
128/32-layout NEON decode (dotprod + scalar fallback); the x86/AVX path is
unchanged. After the fix the CPU forward pass is bit-identical to the
Metal/dequant path (modulo the expected int8-vs-f16 activation difference) and
generation is coherent. Verified on a Jetson Orin Nano (aarch64).

Fixes microsoft#585

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@EUPHEMEME

Copy link
Copy Markdown
Author

Rebased onto current main (0b341e5). The rebase is clean — this patch only touches src/ggml-bitnet-mad.cpp, and the 14 commits since the original base never touched that file, so the diff is byte-identical to what Matheus C. Feitosa (@matheus-c-feitosa) verified.

This PR is not superseded by the submodule move. Matheus C. Feitosa (@matheus-c-feitosa) is right that the cmake build now runs the I2_S kernels from the submodule's ggml/src/ggml-cpu/quants.c, and the layout fix for that path is isHuangXin/llama.cpp#2 (a thorough, differentially-verified fix — worth landing). But the two are complementary, not redundant: at the currently pinned submodule (390c307), Package.swift:18 still compiles spm-sources/ggml-bitnet-mad.cpp, so the Swift-package build — the iOS/macOS on-device path — still runs exactly the kernel this PR corrects. Shipping one without the other leaves ARM wrong on whichever path you didn't patch.

Runtime confirmation, a second architecture. Matheus C. Feitosa (@matheus-c-feitosa) verified on aarch64 Linux (Neoverse-N1); adding arm64 Darwin. CPU-only build (no GPU offload) with this fix, official BitNet-b1.58-2B-4T ggml-model-i2_s.gguf, greedy:

The capital of France is Paris. Paris is a city that is known for its rich history, culture,

Identical to the Metal path on the same machine. (On Apple Silicon Metal masks the bug by dequantizing first, as the PR description notes — this is the CPU NEON path, where the bug is real.)

One caveat for maintainers, with a control. Current main does not build on macOS arm64 at its own pinned submodule (390c307), independent of this PR — two upstream breakages I hit while retesting:

  • ggml/src/ggml-cpu/ggml-cpu.c:1495: use of undeclared identifier 'src1_cont' (compile error in the I2_S fast path).
  • libggml-base: Undefined symbols: _dequantize_row_i2_s, _quantize_i2_s — referenced from ggml.c (ggml-base) but defined in ggml-cpu/quants.c (ggml-cpu); a library-layering break. Control: unmodified main with no patch fails identically, so it is not introduced here.

That's why the runtime receipt above is from a build at this PR's base pin (where it compiles) rather than a fresh build on main — the patch content is identical either way. Happy to extend the ARM QK_I2_S fix to the two submodule sites as well (quants.c, ggml-cpu-i2s.c) if you'd rather have it in one place, though #2 already covers the quants.c layout.

Analysis assisted by Claude Code; all commands were executed and all output checked on the machine.

@matheus-c-feitosa

Copy link
Copy Markdown

Thanks for the rebase and the Darwin run. You're right, and I should have been more precise: "no longer compiles this file" only holds for the CMake build that setup_env.py runs. At 390c307, Package.swift:18 lists spm-sources/ggml-bitnet-mad.cpp, which is a symlink to ../../../src/ggml-bitnet-mad.cpp, so the SwiftPM build still compiles exactly the file this PR fixes. I also compared the rebased diff with the head I tested (a92dbc3): the changed lines are identical. I don't have a Mac, so I haven't tested the SwiftPM path myself.

(Assisted by Claude Code.)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants