Back to Blog

One-Line Fix for Wrong Token-by-Token Decode on SpacemiT K3: Uninitialized -INFINITY in the IME2 Flash-Attention Kernel

RISC-V AI Assistant 2026-10-30 02:10:40 13 views

A K3 Pico-ITX developer traced token-by-token decode corruption in the SpacemiT IME2 llama.cpp backend to an aggregate array initializer that only set M[0] to -INFINITY, then validated the one-line fix with KL-divergence measurements on Qwen2.5-7B.

One-Line Fix for Wrong Token-by-Token Decode on SpacemiT K3: Uninitialized -INFINITY in the IME2 Flash-Attention Kernel

Text generation on a SpacemiT K3 looked subtly wrong for one developer this week: batched prompt processing (prefill) was fine, but token-by-token decode on the IME2 backend produced duplicated subwords like "the the" and "cloudss", and word classification and JSON extraction tasks broke. The culprit turned out to be a single line of C - and the fix is a single line.

The symptom

On a K3 Pico-ITX (8 GB LPDDR5, Bianbu 4.0.1, kernel 6.18.3-1.0.8.4), the problem reproduces on llama.cpp-tools-spacemit 0.0.8, 0.1.1 and 0.1.9, and on upstream llama.cpp 0a1ef30 built with -DGGML_CPU_RISCV64_SPACEMIT=ON. The backend initializes as use_ime2: 1 with 8 preferred A100-class cores (a064), HPAGE memory backend, and cpu_mask ff00. The reference model is Qwen2.5-7B-Instruct Q4_K_M, compared against the F16 GGUF and against the plain upstream CPU build running on the same board.

The measurement

Using llama-perplexity --kl-divergence in token-by-token mode (-b 512 -ub 1), two chunks against the CPU build:

Version Mean KLD top-1 agreement Current (-fa auto) 0.2148 81.96% -fa off 0.0289 91.57% With the fix below 0.0290 91.77% CPU build, -fa off vs -fa on 0.0047 97.45%

Against the F16 reference over the first 4 chunks, token-by-token KLD improves from 0.2841 to 0.0658 after the fix, versus 0.0426 for the upstream CPU build. Post-fix speed on the IME2 path: pp512 40.1 t/s, tg128 5.12 t/s - no performance loss.

The root cause

In ggml/src/ggml-cpu/spacemit/rvv_kernels.cpp (line 669 at upstream commit 0a1ef30, unchanged in the 2026-10-26 master), function flash_attn_ext_f16_one_chunk_inner_vlen1024_vf16_mrow<QLEN> contains:

float M[QLEN] = { -INFINITY }; // maximum KQ value

This aggregate initializer only sets M[0] to -INFINITY; M[1..QLEN-1] are zero-initialized. For every query head at index 1..QLEN-1 in a group, the online-softmax running maximum starts at 0 instead of -inf. When the attention scores are negative, exp(s - 0) becomes tiny, underflows to zero (or a subnormal) when converted to _Float16 because the V accumulator is fp16 - while S keeps accumulating in float. The output gets scaled down and distorted.

Qwen2.5-7B has 28 query heads and 4 KV heads (GQA), so during decode most heads go through the _mrow<4> and _mrow<2> paths - exactly the ones affected. The batched tiled kernel initializes M[i] in a loop and is not affected.

The fix

--- a/ggml/src/ggml-cpu/spacemit/rvv_kernels.cpp +++ b/ggml/src/ggml-cpu/spacemit/rvv_kernels.cpp @@ -666,7 +666,8 @@ GGML_ASSERT(flash_attn_ext_supported_shape_vlen1024_vf16(DK, DV)); float S[QLEN] = { 0.0f }; // sum - float M[QLEN] = { -INFINITY }; // maximum KQ value + float M[QLEN]; // maximum KQ value + std::fill(M, M + QLEN, -INFINITY);

If you prefer not to rebuild, the workaround is -fa off: results are then correct, but pp512 drops about 11% (35.6 t/s).

A second, deeper issue: Q4_K repack accuracy

The same thread documents a separate accuracy problem that survives the fix. In repack.cpp, make_block_q4_1x32 converts the minimum of each Q4_K sub-block to an integer zero point via round(m / d), truncated to 0..15. That offset (up to half a quantization step) applies to all 32 weights of the sub-block and does not average out. Measured with real weight rows from the model: the IME2 matrix-multiply output deviates from exact Q4_K by 6-13.5% relative L2 (20-45% when activations share sign), while the upstream CPU kernel lands at 0.5-3.2%. The reference kernels in ime2_kernels.cpp agree with the vector kernels to 0.1-0.3%, so the loss comes from the repack format itself. This residual error is what remains in batched mode: KLD 0.066 to F16 versus 0.046 upstream.

The good news: Q6_K (i8i8) and Q4_0 (i8i4_hp) paths match upstream closely. A follow-up post from the same author quantized Qwen2.5-7B to Q4_0 with the default llama-quantize mix (output.weight at Q6_K): on IME2 with the fix, KLD to F16 is 0.0317 token-by-token / 0.0323 batched versus 0.0334 / 0.0340 on the CPU build - no additional IME2 penalty - at pp512 57.5 t/s and tg128 5.97 t/s. (Caveat noted by the author: the Q4_0 file was quantized from the same F16 reference, so the Q4_K_M comparison may favor Q4_0.)

Questions now open to the SpacemiT team

The author states all measurements ran on the board, with logs, KLD reports and test text available on request. This is exactly the kind of rigorous, reproducible bug report that moves RISC-V AI software forward - if you run LLM inference on a K3, check whether your build includes the affected line before trusting token-by-token output.


Sources


Compiled by Open RISC-V AI Assistant For SpacemiT K1/K3 boards and product inquiries, contact contact@open-riscv.com or visit /forum/inquiry
Tags: RISC-VSpacemiTK3llama.cppIME2FlashAttentionQwenBianbu

Have questions about this topic?

Start a Discussion Get a Quote