Skip to content

CPU: ik_llama.cpp GEMM kernels for K-quant and I-quant weights on ARM (Q2_K/Q3_K DiTs 1.5x, IQ text encoders up to 2.6x) - #48

Open
danielhanchen wants to merge 9 commits into
cpu-sgemm-f16-bf16from
cpu-iqk-gemm
Open

danielhanchen wants to merge 9 commits into
cpu-sgemm-f16-bf16from
cpu-iqk-gemm

Conversation

@danielhanchen

@danielhanchen danielhanchen commented Nov 5, 2026 •

Copy link
Copy Markdown
Member

Stacked on #47 (cpu-sgemm-f16-bf16, itself on #42). Adds one ggml patch, scripts/unsloth/ggml-patches/0018-ggml-cpu-iqk-gemm.patch. No sd.cpp source change and no ggml submodule bump.

Problem

On ARM CPUs (Linux aarch64, Apple Silicon when a matmul stays on the CPU backend), every K-quant and I-quant weight without a repack runs ggml's vec_dot, one row times one column per call: Q2_K, Q3_K, IQ4_XS, IQ3_S / IQ3_XXS, IQ2_S / IQ2_XS / IQ2_XXS, IQ1_S / IQ1_M, plus Q4_K / Q5_K / Q6_K whenever they are not repacked. x86 does not have this problem: the pinned ggml already has an exact tiled K-quant / I-quant GEMM there (src/ggml-cpu/tiled/) for ne11 >= 8.

Fix

Port the GEMM kernels of ik_llama.cpp (commit 009dce19, MIT, Iwan Kawrakow) into src/ggml-cpu/iqk/, for Q2_K, Q3_K, Q4_K, Q5_K, Q6_K, IQ4_XS, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ1_S and IQ1_M weights times the Q8_K activations ggml already makes:

  • Only ik's kernels that read the upstream block layouts and form the same integer products as vec_dot are ported. Only the float accumulation order differs. ik's on-the-fly re-quantization of weights to Q8_K_R8 is not ported (it rounds the weights a second time).
  • NEON with dotprod, AVX2 and AVX512 (F/BW/DQ/VL/VNNI) kernels. Variants without them, and MSVC builds, compile to a stub that returns false.
  • The hook sits in ggml_compute_forward_mul_mat after src1 is converted to Q8_K, next to the llamafile hook. Repacked and BLAS-routed weights never reach it.
  • Every ported file keeps ik's copyright line, the MIT permission notice and the source commit.

Defaults and switches (read once in ggml_cpu_init):

ARM (NEON + dotprod) x86
default on by batch size (src1 rows): 2-7 rows iqk, 8 rows iqk except iq2_xxs / iq3_s, 9-12 rows iqk for q2_K..q6_K, otherwise tiled (from 8 rows) or vec_dot (1 row); see below
GGML_CPU_IQK=0 off off
GGML_CPU_IQK=1 on on, after the tiled path (every row count below 8, including 1)
GGML_CPU_IQK=2 on on for every row count, ahead of the tiled path

GGML_CPU_IQK_TYPES=q2_K,iq2_xxs,... narrows the types, GGML_CPU_IQK_DEBUG=1 prints the enabled types and the first call per type.

Evidence

Speed, ARM: GitHub ubuntu-24.04-arm (Neoverse N2, 4 cores, armv8.6_2 variant), 2 runners

Same runner, same binary for head and GGML_CPU_IQK=0, 3 reps, median. The kill switch renders the base pixels exactly and matches base time, so it is the A/A.

Run base this PR speedup
Z-Image Turbo Q3_K_M, 512 px, sampling s/step 143.4 / 143.0 97.7 / 97.3 1.47x
Z-Image Turbo Q2_K, 512 px, sampling s/step 153.7 / 154.5 98.2 / 99.8 1.56x
klein text encoder Qwen3-4B UD-IQ2_XXS, s 36.1 / 36.1 14.0 / 14.0 2.57x
klein text encoder Qwen3-4B IQ4_XS, s 17.7 / 17.6 12.7 / 12.8 1.39x

Kernel micro-benchmark (mul_mat, K x N x M = 3840 x 3840 x 1024 for K-quants, 2560 x 9728 x 512 for I-quants, plain buffers, same runners): q2_K 2.18x, q3_K 1.80x, q4_K 1.21x, q5_K 2.29x, q6_K 1.24x, iq4_xs 1.46x, iq3_s 3.24x, iq3_xxs 2.95x, iq2_s 2.18x, iq2_xs 2.14x, iq2_xxs 2.80x, iq1_s 2.82x, iq1_m 2.67x. Both runners agree to within 1%.

Apple M1 (GitHub macos-15 VM, 3 threads): test-backend-ops perf iqk on vs off 1.10x (q4_K) to 2.93x (iq3_xxs). In the default Mac setup Accelerate BLAS still takes the large non-repacked matmuls, so Mac render times are unchanged; with SD_CPU_BLAS=0 the IQ text encoders run 1.3-1.8x faster than through Accelerate.

Speed, x86: small batches only

GitHub runners (Zen 3, 2x Zen 4, Zen 5, Xeon 8573C) and Colab (Broadwell, Zen 2), same-machine A/B. The tiled path already runs these types 1.5-25x faster than vec_dot on large batches. Running the iqk kernels ahead of it for every batch (GGML_CPU_IQK=2) was 5-10% slower on Zen 3, 2-4x slower on the AVX512 parts, and a tie within the A/A spread on Broadwell / Zen 2. Example, Zen 4: Z-Image Q5_K_M sampling 26.0 s (tiled) vs 44.3 s (iqk first).

Small batches are different. test-backend-ops perf -o MUL_MAT (m 4096, k 14336, all 13 types) on Zen 3 (EPYC 7763, three runners), Zen 4 (EPYC 9V74), Zen 5 (EPYC 9V45), Xeon 8370C / 8573C and Core i7-8700B, iqk vs the path that runs otherwise:

src1 rows iqk vs other path x86 default
1 0.63-1.33x vs vec_dot (iq1_m / iq2_s / iq2_xxs slower everywhere) vec_dot
2-7 1.05-3.2x vs vec_dot (one 0.96x: q6_K at 2 rows on Zen 4) iqk
8 1.02-2.9x vs tiled; iq2_xxs 0.84x and iq3_s 0.91x on Zen 5 iqk, not iq2_xxs / iq3_s
12 K-quants 1.01-2.3x vs tiled; I-quants 0.54-1.38x iqk for q2_K..q6_K
16 and up 0.16-1.48x vs tiled, mostly below 1 tiled

These are the modulation / embedding matmuls (one or two rows with batched CFG) and short prompts; sampling of SD1.5 Q3_K at 256 px is unchanged by it (20.95 vs 20.88 s on Zen 3). The pixels move only by accumulation order, as on ARM.

Correctness

  • test-backend-ops -o MUL_MAT / -o GET_ROWS on CPU, with a test-only harness change that loads the scalar x64 variant as the reference: haswell and sapphirerapids variants 2510/2510 MUL_MAT, 240/240 GET_ROWS (GGML_CPU_IQK=2, so the iqk kernels are the ones under test).
  • Apple M1, head build: CPU MUL_MAT 1877/1877, GET_ROWS 171/171; Metal vs CPU MUL_MAT passes for all 13 types.
  • An fp64 oracle (quantized random weights, gaussian / heavy-tailed / constant activations, K up to 12288, ne11 1..1024, broadcast): all 13 types pass on haswell, skylakex, alderlake, zen4, sapphirerapids, NEON under qemu and the real M1. The largest iqk vs vec_dot difference is 3.6e-5 of max|ref| on x86 and 1.8e-6 on the M1, and the error against fp64 is the same as vec_dot's. An engagement counter confirms the kernels ran; with GGML_CPU_IQK=0 it stays at 0.
  • The release CMake (GGML_BACKEND_DL, all CPU variants) builds on x86 (all 14 variants) and with GCC 14 on aarch64. No static initializers in the ported code (dlopen runs them for every variant while ggml scores CPUs).

Image accuracy on ARM (base vs this PR vs a BF16 reference, every shipped K quant and the I-quant text encoders): see the comment below.

Not covered

  • Windows: MSVC builds get the stub, so Windows CPU users keep vec_dot (ARM64 and x64).
  • No i8mm-specific kernels; a real M-series Mac with more than 3 cores and RAM above 7 GB was not available.
  • On Mac the large non-repacked matmuls stay on Accelerate. Routing IQ weights to these kernels instead is a separate change.

…tches/0018)

Ports the classic iqk GEMM kernels of ik_llama.cpp (commit 009dce19, MIT, Iwan Kawrakow) into ggml-cpu as
src/ggml-cpu/iqk/: Q2_K, Q3_K, Q4_K, Q5_K, Q6_K, IQ4_XS, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ1_S and IQ1_M
weights times Q8_K activations, for AVX2, AVX512 (F/BW/DQ/VL/VNNI) and NEON with dotprod. They read the upstream
block layouts and the Q8_K rows ggml already makes, form the same integer products as vec_dot and only accumulate
the floats in another order; ik's re-quantizing Q8_K_R8 conversion path is not ported. Several activation rows
share each weight block decode (up to 8), instead of one row by one column per vec_dot call.

Hooked in ggml_compute_forward_mul_mat after src1 is converted to Q8_K, in place of the vec_dot loop. On x86 the
tiled K-quant path keeps taking ne11 >= 8 first; on ARM, which has no tiled path, these kernels take every ne11.
Variants without AVX2 (or NEON dotprod) and MSVC builds compile a stub and keep vec_dot. Repacked weights do not
reach this code. GGML_CPU_IQK=0 disables, GGML_CPU_IQK=2 runs them ahead of the tiled path,
GGML_CPU_IQK_TYPES=q2_K,iq2_xs,... restricts the types, GGML_CPU_IQK_DEBUG=1 logs the first use per type.
@danielhanchen danielhanchen self-assigned this Nov 5, 2026
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Nov 5, 2026 •

Copy link
Copy Markdown
Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-11-07T08:48:02.991539Z 2f789a4 New commits
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Keep it up!

Reviewed commit: 0ad6060ce4

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@danielhanchen

Copy link
Copy Markdown
Member Author

Image accuracy, ARM (GitHub ubuntu-24.04-arm, Neoverse N2), where this PR is on by default. LPIPS (alex), mean of 2 seeds, 256 px, Studio step counts (Z-Image 8, klein 4). Reference = the same build with BF16 weights for the part under test.

Z-Image Turbo DiT base vs BF16 ref this PR vs BF16 ref base vs this PR
Q2_K 0.3301 0.3244 0.0239
Q3_K_M 0.2713 0.2716 0.0077
Q4_K_M 0.1040 0.1040 0.0000
Q5_K_M 0.0253 0.0253 0.0000
Q6_K 0.0946 0.0946 0.0000
klein 4B text encoder (Qwen3-4B) base vs BF16 ref this PR vs BF16 ref base vs this PR
IQ4_XS 0.0551 0.0568 0.0296
UD-IQ3_XXS 0.2120 0.2037 0.0494
UD-IQ2_XXS 0.3172 0.3070 0.0451
UD-IQ1_M 0.4061 0.3877 0.1221
UD-Q2_K_XL 0.1530 0.1580 0.0361
Q4_K_M 0.0497 0.0497 0.0000
  • A tie against the reference. Per render this PR is closer on 2 of 4 Z-Image Q2_K / Q3_K_M renders and on 7 of the 10 text encoder renders that differ.
  • Q4_K_M / Q5_K_M / Q6_K DiTs and the Q4_K_M text encoder render identical pixels (their weights are repacked on ARM and never reach these kernels).
  • base vs this PR (0.008-0.12) is the usual size of a rounding change in these short renders: on the M1, Accelerate vs vec_dot moves the same text encoders by 0.03-0.22.
  • The large distances of Q2_K / Q3_K_M and IQ1_M / IQ2_XXS / IQ3_XXS to BF16 are the quantization itself, the same in both arms.
  • x86 is unchanged by default (bit-identical to the base).

test-backend-ops MUL_MAT perf (m 4096, k 14336) on Zen 3, Zen 4, Zen 5, Xeon 8370C / 8573C and Core
i7-8700B: iqk is 1.07-3.2x faster than vec_dot for 2-7 src1 rows and 1.15-2.9x faster than the tiled
path at 8 rows (not iq2_xxs / iq3_s, 0.84x / 0.91x on Zen 5), and 1.01-2.3x for K-quants up to 12 rows;
tiled wins from 16 rows up and vec_dot at 1 row. x86 now takes iqk exactly there when GGML_CPU_IQK is
unset; GGML_CPU_IQK=0 / 1 / 2 keep their meaning.

This branch has not been deployed

No deployments
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.

1 participant