diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 35c0551bc..9c009385d 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -2702,10 +2702,13 @@ jobs: echo "CXX=g++-14" >> "$GITHUB_ENV" - name: Display CPU Info uses: bernardladenthin/actions/common/print-host-info@44390bd464395f560f4170d2c4bfb7a6348541b3 # v1.1.0 + # armv8.2-a+dotprod: the plain ARMv8 baseline has no repacked CPU weights at all, so all 24 + # test_moe_repack cases (llama/patches/0022) skipped here; dotprod selects the q4_0_4x4 and + # q4_K_8x4 kernels, the ones an N1-class server runs (the shipped jar's armv8.2 modules) - name: Build libraries run: | mvn -q --no-transfer-progress -f llama/pom.xml compile - .github/build.sh "-DOS_NAME=Linux -DOS_ARCH=aarch64 -DGGML_NATIVE=OFF -DBUILD_TESTING=ON" + .github/build.sh "-DOS_NAME=Linux -DOS_ARCH=aarch64 -DGGML_NATIVE=OFF -DGGML_CPU_ARM_ARCH=armv8.2-a+dotprod+fp16 -DBUILD_TESTING=ON" - name: Verify llama.cpp patches are applied run: .github/verify-patches-applied.sh - name: Run C++ unit tests diff --git a/CHANGELOG.md b/CHANGELOG.md index f4c2f671c..13cccad30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,6 +66,14 @@ from version 5.0.0 onward. Pre-fork releases (`1.x`–`4.2.0`) were authored by An app that used checker-qual's annotations itself through the AAR has to declare it. ### Added +- **Faster prompt processing for MoE models on the CPU** (`llama/patches/0022`): the expert matmuls + over repacked weights batch the tokens of an expert through `gemm` and hand their work out + dynamically, instead of one gemv per token on a fixed thread split. With trained weights + (gpt-oss-20b MXFP4, Qwen3-30B-A3B Q4_K_M, OLMoE-1B-7B Q4_0) prompt processing is +13 % to +66 % + on a Zen 3, a Neoverse-N1 and an IBM z16, generation unchanged + ([`docs/benchmarks/moe-repack-real-weights.md`](docs/benchmarks/moe-repack-real-weights.md)). + `LlamaModel.setCpuMoeBatched(false)` (or `GGML_CPU_REPACK_MMID_BATCHED=0`) restores upstream's + path for a comparison on one's own machine. - **`TextToSpeech.supportsSpeakerReference()`** (and `FfmTextToSpeech.supportsSpeakerReference()`, C API `jllama_tts_supports_speaker_reference()`): whether the loaded pair clones a voice from a speaker reference clip, llama.cpp's `mtmd_helper_cap::audio_gen_spk_ref` (b11901, #31183) -- diff --git a/CLAUDE.md b/CLAUDE.md index 4ba5126d7..6010608ee 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1533,7 +1533,7 @@ Current patches: | `0018-ggml-amx-free-the-aligned-buffer.patch` | **The AMX weight buffer is freed with `free()` but allocated with `_aligned_malloc()` on Windows -- `STATUS_HEAP_CORRUPTION` (0xC0000374) the moment a model is unloaded.** `ggml_backend_amx_buffer_type_alloc_buffer()` (`ggml/src/ggml-cpu/amx/amx.cpp`) allocates with `ggml_aligned_malloc()`, which is `_aligned_malloc()` under `_MSC_VER` (plain clang on a `*-windows-msvc` target defines it too); `ggml_backend_amx_buffer_free_buffer()` called `free()` on it. Fine on Linux (`posix_memalign` + `free`), fatal on Windows, and only there. The patch frees with `ggml_aligned_free(buffer->context, buffer->size)`, exactly what `ggml_backend_cpu_buffer_free_buffer()` does; on Linux that is `free()`, so nothing changes there. **Reached by** the `sapphirerapids` CPU module (the only variant built with `AMX_TILE AMX_INT8`, which `ggml_backend_load_best` picks on Sapphire/Emerald/Granite Rapids; `ggml_amx_init()` is `return true` on `_WIN32`) and a model whose weights have AMX kernels (Q4_0, Q4_1, Q8_0, Q4_K, Q5_K, Q6_K, IQ4_XS -- i.e. every CI model). **Found by the second full run of the modular natives** (38069595101): `Java Tests Windows 2025 x86_64 (clang)` landed on an `Intel(R) Xeon(R) 6973P-C` runner (Granite Rapids), loaded `ggml-cpu-sapphirerapids.dll`, ran `GrammarExample`'s two generations fine and died with exit code -1073740940 at `LlamaModel.close()`; the first run (38065017595) had the same job green on an AMD EPYC 7763 (no AVX-512, `haswell`). **The Windows runners are a CPU lottery**, so which CPU module a Windows job exercises differs per run; before the modular natives the Windows jar was a static Haswell-level build and this code never ran. **Of the six callers of `ggml_aligned_malloc` in the tree this was the only mismatched pair** (`ggml-backend.cpp`, `ggml-cpu.c` twice, `ggml-openvino.cpp`, `ggml-zdnn.cpp` and `ggml.c` all pair correctly), and `ggml-backend.cpp:2464` is the line the fix copies. **Reproduced deliberately rather than waiting for the lottery again**: on a rented Granite Rapids machine, one CPU module staged at a time with the same model and the DLLs CI itself built, `sapphirerapids` crashed **6 of 6** runs -- including at **one** thread, which rules out a data race -- while `cooperlake`, `icelake` and `haswell` were clean 6 of 6, and on Linux with the same CPU all six exit 0 with byte-identical output. The crash lands *after* the inference prints its timings, in cleanup: the signature of a heap corrupted earlier and detected at `free`. With the patch applied the same matrix on the same machine is **clean in all twelve runs**. **The guard is a configure-time check, not a `jllama_test` case**: the AMX code exists only in that one variant of a `GGML_CPU_ALL_VARIANTS` build and executes only on AMX hardware, which the static test build, the local agent's Zen 3 and the cloud sandboxes lack -- so **`llama/cmake/check-aligned-free-pairing.cmake`** fails the configure on **every** platform when a ggml source calls `ggml_aligned_malloc` and never `ggml_aligned_free` (verified red against the unpatched source, naming `amx.cpp`, and green with this patch). A run-time test cannot replace it: without an AMX CPU there is nothing to observe, which is exactly how this shipped. What exercises it: the Windows x86-64 Java test job and the `windows-x86-64` smoke row whenever GitHub hands them an Intel AMX runner; the patch itself is held by `check-patches.py`, applies at b11639, and the patched file compiles with the variant's flags (GCC 13, `-mamx-tile -mamx-int8 -mavx512vnni ...`). Upstream master still carried the `free()` on 2026-10-10; **not yet filed upstream** (nothing before the release). **Drop-check on every bump, note the inversion:** `grep -n "free(buffer->context)" ggml/src/ggml-cpu/amx/amx.cpp` at the new tag -- a **hit** means the bare `free()` is still there and the patch is still needed; **gone** means upstream fixed it, drop the patch. Measured at b11639 (`a0e6dc738`): one hit, line 46. | | `0020-ggml-alloc-validate-leafs-before-reusing-an-allocation.patch` | **An `--embedding` model that is asked to generate crashes the process in `ggml_gallocr_alloc_graph`** (SIGSEGV on Linux, 0xC0000005 with both Windows toolchains), since upstream #30445 (`cb55ffe6e`, "graph: skip the LM head for pooled embeddings", b11574..b11612). `ggml_gallocr_alloc_graph()` assigns every tensor by **position** from the graph it was reserved with, and `ggml_gallocr_needs_realloc()` validates every node and every node's sources but of the **leafs only the count**. Since #30445 the LM head is expanded by `build_logits()` *after* `build_pooling()`, so an embeddings context reserves a graph that expands `t_embd` first and `t_logits` second, while a generation step on it (the server calls `llama_set_embeddings(ctx, false)` for a slot that samples) expands only `t_logits`, whose first source is `output.weight`: same node and leaf counts, every leaf after it shifted by one. An input leaf then lands on the slot of a weight (`buffer_id -1`) and, with the asserts compiled out, `ggml_vbuffer_tensor_alloc()` indexes `galloc->buffers[-1]`; a leaf landing on another input's slot is not size-checked either. The patch validates the leafs exactly as the nodes are validated, so an invalid leaf makes the allocator reserve again (one buffer) or return false, on which `ggml_backend_sched` reserves the graph itself. **Found by** Publish run 38648257387 (2026-10-14, the first full run since b11558): `ErrorHandlingTest` loads its model with `enableEmbedding()` and posts an infill request to it -- its first test in JUnit's order -- and the fork died in all three Java test jobs (`ggml_gallocr_alloc_graph <- ggml_backend_sched_alloc_graph <- llama_context::process_ubatch <- decode <- server_context_impl::update_slots`). **Reproduced locally without a downloaded model** (a random 2-layer llama GGUF with FIM tokens, context 128): infill crashes with `--embedding` and works without; a debug print showed leaf 31 `out_ids` arriving at a slot with no allocation; reverting `cb55ffe6e` makes it pass, and so does this patch -- with greedy sampling the `--embedding` model then returns the same tokens as the model without, run after run. Upstream master (`a0e6dc738`) still lacks the check; **not yet filed upstream**. **Runnable guard: `src/test/cpp/test_gallocr_leafs.cpp`** (two graphs with the same counts and the leafs swapped). **Drop-check on every bump:** read `ggml_gallocr_needs_realloc()` in `ggml/src/ggml-alloc.c` at the new tag -- once it walks `graph->leafs` itself, drop this patch rather than refreshing it. | | `0021-server-clear-the-slot-cache-when-the-server-adapters-change.patch` | **A slot's prompt cache survived a change of the LoRA adapters it was computed with.** `launch_slot_with_task()` (`tools/server/server-context.cpp`) clears the cache when a request names its own adapters and they differ from the slot's (`lora_should_clear_cache`, which keeps it only for aLoRA), but a request that names none takes `params_base.lora_adapters` in the `else` branch without that check. So the cache survived `SET_LORA` (`LlamaModel.setLoraAdapter(s)`, upstream's `POST /lora-adapters`) and the return to the server's adapters after a request with its own, and the slot answered like neither adapter set, silently. **Measured** (J9 of the bug-hunt brief, b11639, upstream's own LoRA test pair `stories15M_MOE-F16.gguf` + `moe_shakespeare15M.gguf`, greedy, one slot): after `setLoraAdapter(0, 0f)` the warm slot answered ". He saw a picture of a bird ..." where an emptied slot with the same scale answers " was a little girl ..."; only `eraseSlot(0)` restored it (F67). The patch applies the per-request rule in the `else` branch too; equal sets keep the cache as before. **A second hunk starts every slot with the server's adapters in `init()`** (`slot.lora = params_base.lora_adapters;`): `slot.lora` used to start empty, so a never-used slot of a server started with `--lora` read as "computed under no adapters", and its first task logged `0 loras -> 1 loras` and cleared a prompt it already held -- a state restored into it (`SLOT_RESTORE`) or loaded from the host-memory prompt cache by `get_available_slot()`, both computed under the server's adapters (F240). Upstream has the same defect on its HTTP route, so the patch is upstream-shaped; **not yet filed upstream**. Not covered, and said in `setLoraAdapters`' Javadoc: the host-memory prompt cache (`--cache-ram`, on by default) stores slot states without their adapters. **Runnable guard: `src/test/cpp/test_lora_cache.cpp`** (a random llama model and a random LoRA adapter on K, V and the FFN output, written by the test; five cases). **Drop-check on every bump, one per hunk:** read the `else` branch of the "process per-request lora adapters" block in `launch_slot_with_task()` at the new tag -- once it compares `params_base.lora_adapters` with `slot.lora` and clears the prompt (or `SET_LORA` clears the slots), drop that hunk rather than refreshing it; and read the slot loop of `init()` -- once it sets `slot.lora` from `params_base.lora_adapters`, drop that hunk. The patch goes when both have. | -| `0022-ggml-cpu-repack-batch-the-tokens-of-an-expert.patch` | **The expert matmuls of a MoE model over repacked CPU weights ran one token at a time on a static thread split.** `forward_mul_mat_id()` (`ggml/src/ggml-cpu/repack.cpp`) -- the path of every MoE model whose experts are repacked at load: MXFP4 (gpt-oss), Q4_0, IQ4_NL and Q4_K on AVX2, their interleaved forms on ARM -- multiplied each token routed to an expert with its own gemv, although the gemm kernel for the type exists, and gave thread `ith` a fixed row range of every expert, so each expert matmul waited for the slowest thread. Measured with a random-weight gpt-oss-20b in ggml-org's MXFP4 layout on the `alderlake` module (the Core Ultra's), 4 threads: 72 % of prompt time was that gemv loop, 13.9 % of generation time the barrier wait after it ([`docs/benchmarks/moe-mxfp4-cpu-alderlake.md`](docs/benchmarks/moe-mxfp4-cpu-alderlake.md)). The patch adds a batched path -- (expert, column chunk) units handed out through `ggml_threadpool_chunk_add`, every group of up to 16 of an expert's tokens quantized interleaved (`ggml_quantize_mat_t`, as `forward_mul_mat` does) and multiplied with `gemm`, a rest of fewer than four by gemv -- and **keeps upstream's path behind a process-wide switch**: `ggml_backend_cpu_set/get_repack_mmid_batched` (`ggml-cpu.h`, and in the CPU backend's `get_proc_address` table, which is how libjllama reaches it in a `GGML_BACKEND_DL` build), on by default, `GGML_CPU_REPACK_MMID_BATCHED=0` in the environment starts it off. Thread 0 reads it once per op and publishes the choice in the work buffer, so a flip during a compute cannot split the threads; `work_size()` reserves the batched path's scratch either way (~420 KB per thread for gpt-oss). Exposed as `LlamaModel.setCpuMoeBatched` / `isCpuMoeBatched`, `jllama_set/get_cpu_moe_batched` and the FFM `Jllama`: the switch exists so a user can measure both paths on their own hardware -- the P-core/E-core effect the static split suggests is unmeasured (TODO.md). **Measured, one binary, the environment variable flipped:** prompt processing +47 % (42.63 -> 62.76 t/s), token generation +26 % (7.90 -> 9.97), logits equal to rounding (`llama-perplexity --kl-divergence`: mean KLD 0.000000, max 0.000002, same top token 100 %). **Runnable guard: `src/test/cpp/test_moe_repack.cpp`** (both paths against ggml's generic `MUL_MAT_ID` over the same weights; the switch resolved through `get_proc_address`). Upstream-shaped without the switch; **not yet filed upstream**. **Drop-check on every bump:** read `forward_mul_mat_id()` in `ggml/src/ggml-cpu/repack.cpp` at the new tag -- once it calls `gemm` for an expert's batched rows and hands its work out through `ggml_threadpool_chunk_add`, drop this patch rather than refreshing it; the switch goes with it, and `setCpuMoeBatched` then throws (`LlamaException`, no entry point), so the Java, C and FFM API needs the decision whether to remove it or keep it as a no-op. | +| `0022-ggml-cpu-repack-batch-the-tokens-of-an-expert.patch` | **The expert matmuls of a MoE model over repacked CPU weights ran one token at a time on a static thread split.** `forward_mul_mat_id()` (`ggml/src/ggml-cpu/repack.cpp`) -- the path of every MoE model whose experts are repacked at load: MXFP4 (gpt-oss), Q4_0, IQ4_NL and Q4_K on AVX2, their interleaved forms on ARM -- multiplied each token routed to an expert with its own gemv, although the gemm kernel for the type exists, and gave thread `ith` a fixed row range of every expert, so each expert matmul waited for the slowest thread. Measured with a random-weight gpt-oss-20b in ggml-org's MXFP4 layout on the `alderlake` module (the Core Ultra's), 4 threads: 72 % of prompt time was that gemv loop, 13.9 % of generation time the barrier wait after it ([`docs/benchmarks/moe-mxfp4-cpu-alderlake.md`](docs/benchmarks/moe-mxfp4-cpu-alderlake.md)). The patch adds a batched path -- every group of four tokens routed to an expert gathered and quantized interleaved (`ggml_quantize_mat_t`, as `forward_mul_mat` does) **once**, into a buffer shared by all threads, then (expert, column chunk) units handed out through `ggml_threadpool_chunk_add`, each multiplying up to 16 shared rows at a time with `gemm`, a rest of fewer than four by gemv -- and **keeps upstream's path behind a process-wide switch**: `ggml_backend_cpu_set/get_repack_mmid_batched` (`ggml-cpu.h`, and in the CPU backend's `get_proc_address` table, which is how libjllama reaches it in a `GGML_BACKEND_DL` build), on by default, `GGML_CPU_REPACK_MMID_BATCHED=0` in the environment starts it off. Thread 0 reads it once per op and publishes the choice in the work buffer, so a flip during a compute cannot split the threads; `work_size()` reserves the batched path's shared rows (one per routed token slot, 6.3 MB for gpt-oss at a 512-token batch) and per-thread scratch either way. **Quantizing once is load-bearing on ARM:** the first version quantized an expert's rows again in every column chunk, and with no SIMD Q8_K quantizer on ARM that made Q4_K prompt processing on a Neoverse-N1 18 % *slower* than upstream at 32 threads -- found with real weights ([`docs/benchmarks/moe-repack-real-weights.md`](docs/benchmarks/moe-repack-real-weights.md)), where the final patch gains +13 % to +66 % in prompt processing on Zen 3, the N1 and a z16 (Q4_0 only, the one type s390x repacks) for MXFP4, Q4_0 and Q4_K, generation unchanged; KLD 0 for MXFP4/Q4_0, 0.0075 for Q4_K, where gemm and gemv round differently and upstream's own repack path is as far from `--no-repack`. Exposed as `LlamaModel.setCpuMoeBatched` / `isCpuMoeBatched`, `jllama_set/get_cpu_moe_batched` and the FFM `Jllama`: the switch exists so a user can measure both paths on their own hardware -- the P-core/E-core effect the static split suggests is unmeasured (TODO.md). **Measured, one binary, the environment variable flipped:** prompt processing +47 % (42.63 -> 62.76 t/s), token generation +26 % (7.90 -> 9.97), logits equal to rounding (`llama-perplexity --kl-divergence`: mean KLD 0.000000, max 0.000002, same top token 100 %). **Runnable guard: `src/test/cpp/test_moe_repack.cpp`** (both paths against ggml's generic `MUL_MAT_ID` over the same weights; the switch resolved through `get_proc_address`). Upstream-shaped without the switch; **not yet filed upstream**. **Drop-check on every bump:** read `forward_mul_mat_id()` in `ggml/src/ggml-cpu/repack.cpp` at the new tag -- once it calls `gemm` for an expert's batched rows and hands its work out through `ggml_threadpool_chunk_add`, drop this patch rather than refreshing it; the switch goes with it, and `setCpuMoeBatched` then throws (`LlamaException`, no entry point), so the Java, C and FFM API needs the decision whether to remove it or keep it as a no-op. | | `0023-ggml-cpu-s390x-variants-build-with-clang-and-a-z13-baseline.patch` | **The IBM Z CPU variants (`GGML_CPU_ALL_VARIANTS` on s390x) did not build with clang, and had no module below z15.** (1) `ggml-cpu-impl.h` and `llamafile/sgemm.h` both include `` under `__VXE__ \|\| __VXE2__`; clang's header has no include guard (GCC's has `_VECINTRIN_H`), so `ggml-cpu.c`, which includes both, failed with "redefinition of 'vec_extract'". The include in `sgemm.h` came with ggml-org/llama.cpp#14273 and is not needed: the header declares no vector type, and both of its includers include `ggml-cpu-impl.h` -- `sgemm.cpp` after it, where `__VXE__` is not defined yet (`ggml-cpu-impl.h` defines it from `__VEC__`), so it never fired there. The patch removes it. (2) Upstream's list starts at `z15`, so a z13/z14 would have found no CPU backend at all; the patch adds `ggml-cpu-z13` (scalar, `-march=z13`, score 1), so the jar keeps the z13 floor the static library had. Measured on cfarm191 (z16, b11901, Qwen3-0.6B Q4_K_M big-endian, 4 threads): `z13` 19.7 / 13.6 t/s, `z15` 90.0-95.5 / 57.3-58.3, `z16` 82.6-88.3 / 51.9-53.9 (prompt / generation); Q4_K_M against a Q8_0 reference on z13: mean KLD 0.0777 on `z13`, 0.0746 on `z16`. **Guard:** `build-linux-s390x` builds the variants with clang 18 -- without part (1) the build fails -- and checks that `libggml-cpu-z13/z15/z16.so` all exist, which is what fails without part (2). Not yet filed upstream. **Drop-check on every bump:** `grep -n vecintrin ggml/src/ggml-cpu/llamafile/sgemm.h` (a hit: part 1 still needed) and the s390x block of `ggml_add_cpu_backend_variant` calls in `ggml/src/CMakeLists.txt` (no `z13`/scalar entry: part 2 still needed). Measured at b11901: both still needed. | **Dropped patches** (`0004`, `0005`, `0009`, `0010`, `0011`, `0012`, `0013`, `0016`, `0019` -- upstream fixed @@ -2821,7 +2821,7 @@ ctest --test-dir build --output-on-failure -R "ResultsToJson" | `src/test/cpp/test_native_server_shutdown.cpp` | 3 | **The runnable guard for the shutdown-handler guard in `patches/0006`/`0007`.** Runs the real `llama_server()` in router mode over an empty `--models-dir` on an ephemeral loopback port (no model, no worker) and stops it the way `native_server.cpp` does. Pins a clean stop (exit code 0), that `llama_server_request_shutdown()` **after** the server returned is a no-op (the deterministic form of the CI `SIGSEGV` — it crashed every run before the fix), and that requests hammered from another thread while the server tears down are safe (5 rounds). Compiled on non-Android only, like `server.cpp` itself. | | `src/test/cpp/test_gallocr_leafs.cpp` | 1 | **The runnable guard for `patches/0020`.** Two graphs with the same two nodes in the same order (`n1 = add(x, x)`, `n2 = mul(w, n1)`) and only the leafs moved: one expands `n1` before `n2` and meets the input first, the other expands only `n2`, whose first source is the weight -- the shape of the llama.cpp case, where every node and node source still matches its slot, so the old check passed it. Reserved with the first, the second must get an allocation for the input inside the allocator's buffer and compute the right product on it. No model, every platform. **Verified both ways:** with the fix removed from the source tree the test dies with SIGSEGV (rc 139); a first version that swapped the operands of one `add` passed without the fix, because a swapped *node source* was already caught -- a guard has to move only the leafs. | | `src/test/cpp/test_lora_cache.cpp` | 5 | **The runnable guard for `patches/0021`.** Writes a random llama model (the shared `tiny_llama_gguf.hpp`, also used by `test_c_api.cpp`) and a random LoRA adapter on K, V and the FFN output of every layer, loads them with one slot and `--cache-ram 0`, and answers the same prompt greedily. After `jllama_set_lora_adapters` (scale 1 -> 0 and back) and after a request with its own adapters, the warm slot must answer exactly like an emptied one. One case guards the others against passing vacuously: the adapter changes the answer, and a repeated request reuses the cache (`timings.cache_n > 0`; not `tokens_cached`, which counts every token the slot holds and is positive after any request). After each adapter change the next request must take nothing from the cache (`timings.cache_n == 0`), so those three cases are red without the patch whether or not the random model happens to answer differently. **Verified both ways:** without the patch 3 of 4 red (the warm slot answered from the old scales), with it 4 of 4 green in three runs; no model, every platform. A second fixture (`LoraCacheFreshSlot`: two slots, `--slot-save-path`) saves slot 0 (by its file name inside that directory) after a request and restores it into the never-used slot 1, whose request must take all of the prompt but at most two tokens from it and answer like slot 0 -- red before the `init()` hunk (`Expected: (prompt_n) <= (2), actual: 77 vs 2`), green with it (F240). | -| `src/test/cpp/test_moe_repack.cpp` | 2 | **The runnable guard for `patches/0022`.** `MUL_MAT_ID` over weights in the `CPU_REPACK` buffer -- MXFP4, Q4_0, IQ4_NL and Q4_K, random, quantized with `ggml_quantize_chunk` -- against ggml's generic `MUL_MAT_ID` over the same bytes in plain memory (independent code, the same activation quantization), for both settings of the switch on 1, 3 and 4 threads (3 does not divide the 9 column blocks, 1 leaves a unit past the last column). Routings: experts with 0, 1, 3, 4, 15 and 17 tokens (the gemv rest alone, one gemm group, a full 16-row group plus a rest), the same with one slot of -1 (its row stays zero), and a single token; each with one input row per token (gate/up) and one per slot (down): 24 instances. A type the running CPU does not repack is skipped, not passed. The second test resolves the switch through the CPU backend's `get_proc_address`, as libjllama does, and round-trips it, so a bump that drops the patch reds this file. | +| `src/test/cpp/test_moe_repack.cpp` | 2 | **The runnable guard for `patches/0022`.** `MUL_MAT_ID` over weights in the `CPU_REPACK` buffer -- MXFP4, Q4_0, IQ4_NL and Q4_K, random, quantized with `ggml_quantize_chunk` -- against ggml's generic `MUL_MAT_ID` over the same bytes in plain memory (independent code, the same activation quantization), for both settings of the switch on 1, 3 and 4 threads (3 does not divide the 9 column blocks, 1 leaves a unit past the last column). Routings: experts with 0, 1, 3, 4, 15 and 17 tokens (the gemv rest alone, one gemm group, a full 16-row group plus a rest), the same with one slot of -1 (its row stays zero), and a single token; each with one input row per token (gate/up) and one per slot (down): 24 instances. A type the running CPU does not repack is skipped, not passed -- which is why `test-cpp-linux-aarch64` builds with `GGML_CPU_ARM_ARCH=armv8.2-a+dotprod+fp16`: on plain ARMv8 nothing is repacked and all 24 cases skipped (they still do in the Windows arm64 and s390x test jobs). The second test resolves the switch through the CPU backend's `get_proc_address`, as libjllama does, and round-trips it, so a bump that drops the patch reds this file. | **Current total: 681 tests** (`TEST`/`TEST_F`/`TEST_P` definitions; ctest counts 725 cases, because the parameterized `CApiVocabOnly` case expands to 22 instances and `MoeRepack` to 24). diff --git a/README.md b/README.md index a876e819c..b97b75510 100644 --- a/README.md +++ b/README.md @@ -1688,6 +1688,9 @@ core takes fewer of them. Upstream llama.cpp multiplies one token at a time and fixed share. The results are the same up to rounding; measured on gpt-oss-20b in a sandbox, prompt processing was 47 % faster and token generation 26 % ([`docs/benchmarks/moe-mxfp4-cpu-alderlake.md`](docs/benchmarks/moe-mxfp4-cpu-alderlake.md)). +With trained weights (gpt-oss-20b, Qwen3-30B-A3B, OLMoE) prompt processing was 13 % to 66 % faster +on a Zen 3, a Neoverse-N1 and an IBM z16, token generation unchanged +([`docs/benchmarks/moe-repack-real-weights.md`](docs/benchmarks/moe-repack-real-weights.md)). Both paths stay available, so the difference can be measured on your own machine: diff --git a/TODO.md b/TODO.md index bd3b2601c..74ca06047 100644 --- a/TODO.md +++ b/TODO.md @@ -165,7 +165,7 @@ Kept as a list because each entry names the machine it needs -- the harness is | **OpenVINO on an Intel GPU/NPU** | suitable Intel device, unavailable to these sessions | the 21.7x/336x figures are CPU-plugin fallback only; they do not prove GPU/NPU execution | | the two **SME** aarch64 modules | Linux aarch64 with SME, absent from current CI/cloud machines | macOS's static Metal library does not exercise these modules | | **gpt-oss-20b on both MoE paths of `llama/patches/0022`**, on the **Core Ultra 7 165U** (`-t 6/8/10`) | the laptop; switch with `LlamaModel.setCpuMoeBatched` or `GGML_CPU_REPACK_MMID_BATCHED=0` ([`docs/benchmarks/moe-mxfp4-cpu-alderlake.md`](docs/benchmarks/moe-mxfp4-cpu-alderlake.md)) | upstream's static work split meets real P-, E- and LP-E-cores only there; the sandbox could emulate one slow core for prompt processing (−49 % upstream's path, −24 % batched) but not for generation | -| **`llama/patches/0022` with real weights on every CPU family whose repack path it changes**: gpt-oss-20b (MXFP4) and a Q4_K MoE (e.g. Qwen3-30B-A3B Q4_K_M), both switch positions -- `llama-perplexity --kl-divergence` of the batched path against upstream's, `llama-bench` pp/tg at several thread counts, and once through the jar (`LlamaModel.setCpuMoeBatched`) | x86-64 with AVX-512/AMX and Zen 4 (EC2 or Vultr: the `sapphirerapids` and `zen4` modules); aarch64 with dotprod (cfarm424, Neoverse-N1, 64 cores) and with i8mm/SVE (Graviton 3/4); s390x with VXE (cfarm191, z16 -- ggml repacks only Q4_0 there, so a Q4_0 MoE converted to big-endian, on the vector modules `z15`/`z16`, see "Linux s390x" in CLAUDE.md) | #578 measured random weights in the real shapes, on the `alderlake` module of an AVX-512 host. KLD 0 there shows the batched path is right for that gemm kernel, not for the NEON, i8mm or VXE kernels it calls elsewhere, and random weights route tokens to experts more evenly than a trained model does | +| **`llama/patches/0022` with real weights on the CPU families still missing** -- done on Zen 3 (AVX2, `haswell`), a Neoverse-N1 (dotprod, `armv8.2_2`) and an IBM z16 (VXE, `z16`), with gpt-oss-20b, Qwen3-30B-A3B Q4_K_M and OLMoE Q4_0 ([`docs/benchmarks/moe-repack-real-weights.md`](docs/benchmarks/moe-repack-real-weights.md): +13 % to +66 % prompt processing, generation unchanged, KLD 0 for MXFP4/Q4_0); the same `llama-bench` + `llama-perplexity --kl-divergence` pair, both switch positions, on the rest | x86-64 with AVX-512/AMX and Zen 4 (EC2 or Vultr, paid: the `sapphirerapids` and `zen4` modules); aarch64 with i8mm/SVE (Graviton 3/4, paid: the `q4_K_8x8` / `q4_0_4x8` kernels) | each module calls other gemm kernels; the N1 run found a defect (Q4_K −18 % at 32 threads, fixed) that x86 did not show. The C++ guard (`test_moe_repack`) still skips on Windows arm64 and s390x, where nothing is repacked in the test build | | [`docs/local-test-plan.md`](docs/local-test-plan.md) parts A–D on the **Core Ultra 7 165U** | the laptop the Windows guide (#575) is written for | part A -- which CPU module ggml picks there and what it gains -- is what CI cannot measure; the only report so far is the Zen 3 one (`docs/handover/local-agent-report-b11538-windows.md`). One result file per machine, unmeasured values marked as such | ## Open — jllama-specific diff --git a/docs/benchmarks/README.md b/docs/benchmarks/README.md index 5e891c23f..f3cb3b325 100644 --- a/docs/benchmarks/README.md +++ b/docs/benchmarks/README.md @@ -15,6 +15,7 @@ themselves live in [`../../TODO.md`](../../TODO.md). |---|---| | [`cpu-backend-modules.md`](cpu-backend-modules.md) | What each of the 22 shipped CPU backend modules is worth, per microarchitecture — and whether `ggml_backend_load_best` picks the fastest one | | [`moe-mxfp4-cpu-alderlake.md`](moe-mxfp4-cpu-alderlake.md) | Where the `alderlake` module (Core Ultra) loses time on Qwen3 and gpt-oss-20b, measured with random-weight models in the real shapes -- and what the fix to ggml's MoE path in `llama/patches/0022` is worth (+47 % prompt, +26 % generation; switchable for A/B) | +| [`moe-repack-real-weights.md`](moe-repack-real-weights.md) | `llama/patches/0022` with trained weights (gpt-oss-20b MXFP4, Qwen3-30B-A3B Q4_K_M, OLMoE Q4_0) on Zen 3, a Neoverse-N1 and an IBM z16: +13 % to +66 % prompt processing, generation unchanged, KLD against upstream's path -- and the ARM defect of the first version it found | | [`gpu-module-jars.md`](gpu-module-jars.md) | What happens when a GPU module jar is actually *used* -- two findings that contradict the "a module that cannot help is skipped" premise | | [`arm-fp16-conversion.md`](arm-fp16-conversion.md) | Every shipped ARM binary converts fp32 to fp16 in software, ours and upstream's: 26 of 26 disassembled, 8.9x to 13.3x on the function, and 0.03 % of cycles so inference does not move | | [`upstream-f16c-30270.md`](upstream-f16c-30270.md) | An isolation of one upstream commit: what it does, that it changes nothing measurable here, and a compile-time guard proven red against the defect | diff --git a/docs/benchmarks/moe-repack-real-weights.md b/docs/benchmarks/moe-repack-real-weights.md new file mode 100644 index 000000000..c42dad27e --- /dev/null +++ b/docs/benchmarks/moe-repack-real-weights.md @@ -0,0 +1,167 @@ + + +# Patch 0022 with real weights + +**The question.** [`llama/patches/0022`](../patches.md#0022-ggml-cpu-repack-batch-the-tokens-of-an-expertpatch) +batches the tokens of an expert in the CPU backend's repacked `MUL_MAT_ID` and hands its work out +dynamically. It was measured on random weights in gpt-oss-20b's shape, on one x86 module +([`moe-mxfp4-cpu-alderlake.md`](moe-mxfp4-cpu-alderlake.md)). Does it hold with trained weights, +whose routers spread tokens unevenly, on the other kernels it calls: ARM's, K-quants? + +**The answer, measured 2026-10-27 on b11901 + `llama/patches/`:** + +1. **Prompt processing gains on every model and machine measured**: +13 % to +66 %, on x86 + (Zen 3), ARM (Neoverse-N1) and s390x (z16). +2. **Token generation is unchanged** within noise. With one token per expert, nothing is batched. +3. **The logits are those of upstream's path for MXFP4 and Q4_0** (KLD 0.000000, same top token + 100 %). **For Q4_K they differ by a KLD of 0.0075**, which is not a defect of the patch: + upstream's own repack path is as far from the unrepacked computation (below). +4. **It found a defect in the first version of the patch.** On a Neoverse-N1, Q4_K prompt + processing was 18 % slower than upstream's path at 32 threads, because every column chunk + quantized the expert's activations again. ARM has no SIMD form of the Q8_K quantizer. Fixed: + each group of four rows is now quantized once. +5. **The ARM C++ tests never ran the patch's guard.** `test-cpp-linux-aarch64` built for plain + ARMv8, where nothing is repacked, so all 24 `test_moe_repack` cases skipped. It builds with + dotprod now; on the N1 all 25 tests pass. + +## Setup + +- **Models** (trained weights from Hugging Face): + - gpt-oss-20b MXFP4 (`ggml-org/gpt-oss-20b-GGUF`, 12.1 GB); + - Qwen3-30B-A3B Q4_K_M (`unsloth/Qwen3-30B-A3B-GGUF`, 18.6 GB; 128 experts, 8 used); + - OLMoE-1B-7B-0125-Instruct Q4_0 (`allenai/OLMoE-1B-7B-0125-Instruct-GGUF`, 3.9 GB; 64 experts, + 8 used). +- **The build:** llama.cpp b11901 with the 14 patches, built as the modular library jars are + (`GGML_BACKEND_DL`, every CPU variant, no OpenMP, clang). ggml picks the module, upstream's + `llama-bench` and `llama-perplexity` measure. +- **The switch:** `GGML_CPU_REPACK_MMID_BATCHED=0` gives upstream's path, `1` the batched one, in + the same binary. +- **Benchmarks:** `llama-bench -p 512 -n 128 -r 3`, off and on interleaved per thread count. +- **KLD:** `llama-perplexity --kl-divergence` on wikitext-2 (20 chunks of 512), the batched path + against upstream's. +- **The machines:** + - **x86:** Ryzen 7 5800H (Zen 3, 8 cores / 16 threads, module `haswell`), Docker on Windows; + downloads were running during part of the first series. + - **ARM:** Neoverse-N1, cfarm424 (64 cores, dotprod, no i8mm, module `armv8.2_2`), a shared + host; some rounds ran next to other users' load. + - **s390x:** IBM z16, cfarm191 (8 CPUs, module `z16`), OLMoE only ([below](#s390x)). + +## Prompt processing (pp512, t/s, off → on) + +| model (type) | machine | threads | upstream | batched (final) | | +|---|---|--:|--:|--:|--:| +| gpt-oss-20b (MXFP4) | Zen 3 | 4 | 28.56 | 37.23 | +30 % | +| | | 8 | 41.68 | 59.58 | +43 % | +| | | 16 (SMT) | 54.44 | 69.80 | +28 % | +| | N1 | 8 | 49.32 | 80.85 | +64 % | +| | | 16 | 98.28 | 152.46 | +55 % | +| | | 32 | 176.97 | 278.70 | +57 % | +| Qwen3-30B-A3B (Q4_K) | Zen 3 | 4 | 39.96 | 53.27 | +33 % | +| | | 8 | 58.32 | 84.65 | +45 % | +| | | 16 (SMT) | 71.72 | 82.27 | +15 % | +| | N1 | 8 | 56.81 | 66.77 | +18 % | +| | | 16 | 105.99 | 121.84 | +15 % | +| | | 32 | 190.66 | 215.69 | +13 % | +| OLMoE-1B-7B (Q4_0) | Zen 3 | 4 | 81.11 | 116.78 | +44 % | +| | | 8 | 138.81 | 189.95 | +37 % | +| | | 16 (SMT) | 204.83 | 236.12 | +15 % | +| | N1 | 8 | 192.28 | 274.53 | +43 % | +| | | 16 | 378.50 | 506.39 | +34 % | +| | | 32 | 705.85 | 989.98 | +40 % | + +Rows with the final patch: Zen 3 at 8 threads, and N1 Qwen3 plus every N1 row at 32 threads. The +other rows were measured with the first version. For MXFP4 and Q4_0 the two versions do the same +work, and the activation quantizers have SIMD forms on both machines. + +**Token generation** (tg128) stays within the spread at every point, e.g. gpt-oss on N1 at 32 +threads 42.60 → 42.10, Qwen3 on Zen 3 at 8 threads 16.42 → 16.51. One N1 run at 32 threads read +38.03 → 31.86 ± 3.76 for Qwen3. Repeated in alternating rounds on that host (load average 33 from +other users), it read 38.00 / 35.10 and 32.20 / 35.11: noise. + +## The defect the real weights found + +The first version of the batched path quantized an expert's rows inside every (expert, column +chunk) unit. With more threads there are more column chunks per expert (`nck`), so the same rows +were quantized up to four times at 32 threads. On the N1, Qwen3-30B (Q4_K): + +| threads | upstream | first version | final | +|--:|--:|--:|--:| +| 8 | 56.66 | 62.22 (+10 %) | 66.77 (+18 %) | +| 16 | 107.62 | 104.95 (−2 %) | 121.84 (+15 %) | +| 32 | 200.91 | 164.04 (**−18 %**) | 215.69 (+13 %) | + +The upstream column of the final run is from a later series on the shared host: 56.81 / 105.99 / +190.66. + +`perf` at 32 threads: upstream's path spends 36.5 % in `ggml_gemv_q4_K_8x4_q8_K` and 13.9 % in +`ggml_gemv_q6_K_8x4_q8_K`. The first version spent 30.3 % in `ggml_gemm_q4_K_8x4_q8_K` and 5.5 % in +the Q6_K gemm, but **20.8 % inside `forward_mul_mat_id` itself** (the inlined +`ggml_quantize_mat_q8_K_4x4_generic`) plus 7.5 % in `memcpy`. x86 has an AVX2 +`ggml_quantize_mat_q8_K_4x8`; ARM has no SIMD Q8_K quantizer at all. MXFP4 and Q4_0 quantize to +Q8_0, which has NEON forms, so only K-quants on ARM showed it. + +The final version gathers and quantizes every group of four rows once, into a buffer the units +share (`work_size` reserves one row per routed token slot: 9.6 MB for Qwen3-30B at a 512-token +batch). The groups are split statically over the threads, then a barrier; both are skipped when no +expert got four tokens. + +## KLD: why Q4_K is not zero + +| model | machine | batched vs upstream: mean KLD | max | same top token | +|---|---|--:|--:|--:| +| gpt-oss-20b (MXFP4) | Zen 3 | 0.000000 | 0.000052 | 100 % | +| | N1 | 0.000000 | 0.000048 | 100 % | +| OLMoE-1B-7B (Q4_0) | Zen 3 | 0.000000 | 0.000063 | 100 % | +| | N1 | 0.000000 | 0.000056 | 100 % | +| Qwen3-30B-A3B (Q4_K) | Zen 3 | 0.007511 | 0.929 | 95.8 % | +| | N1 | 0.007316 | 0.897 | 96.1 % | + +For MXFP4 and Q4_0, gemm and gemv round alike, so batching changes nothing. For Q4_K they round +differently (another summation order of the Q8_K block sums). In an MoE model, a router near a tie +then picks another expert, which turns a rounding difference into a different token now and then. + +**The reference settles which path is "right".** Both paths were compared with the unrepacked +generic computation (`--no-repack`, 8 chunks, N1): + +- upstream's repack path: mean KLD 0.007986, same top token 95.5 %; +- the batched path: mean KLD 0.008496, same top token 95.0 %. + +Both are as far from it as from each other. The final version gives the same KLD as the first, +digit for digit, since it computes the same products. + +## Through the jar + +`LlamaModel.setCpuMoeBatched(false/true)` alternated in one JVM on the same chat request, prompt +cache off, the CI-built 6.1.0-SNAPSHOT `cpu-linux-aarch64` jar (first version of the patch), N1, +16 threads: + +- gpt-oss-20b: prompt 79.4 / 81.1 → 130.7 / 130.6 t/s (+63 %), generation 29.9 → 29.3; +- Qwen3-30B: 82.5 / 84.5 → 82.6 / 82.7 (the first version's Q4_K result at 16 threads). + +The switch reaches the native path through the jar, and the loader picked `armv8.2_2`. + +## s390x + +ggml repacks only Q4_0 on s390x (the VXE `q4_0_4x8` kernels), so OLMoE-1B-7B Q4_0, converted to +big-endian with `gguf_convert_endian.py`, on cfarm191 (IBM z16, 8 CPUs, module `z16`; a shared host, +load average 5 from other users during the run): + +| threads | pp512 upstream | batched | | tg128 upstream | batched | +|--:|--:|--:|--:|--:|--:| +| 2 | 24.34 | 38.97 | +60 % | 8.30 | 8.26 | +| 4 | 45.79 | 76.23 | +66 % | 17.04 | 17.37 | +| 8 | 43.33 | 59.03 | +36 % | 0.31 | 0.33 | + +At 8 threads every CPU was shared with other users' processes, which spinning threads punish: +generation collapsed on both paths alike. KLD of the batched path against upstream's (8 threads, +20 chunks): mean 0.000000, max 0.000059, same top token 99.98 %. + +## Not measured here + +AVX-512/AMX (`sapphirerapids`) and Zen 4 (`zen4`), which need a paid cloud instance; ARM with i8mm +(Graviton 3/4: the `q4_K_8x8` / `q4_0_4x8` kernels); and the hybrid-core case the original +measurement emulated. diff --git a/docs/patches.md b/docs/patches.md index 6c7551889..2511a0ed1 100644 --- a/docs/patches.md +++ b/docs/patches.md @@ -547,15 +547,22 @@ compute cannot send the threads down different paths: (ggml-cpu.h, and in the CPU backend's get_proc_address table for a GGML_BACKEND_DL build). On by default; GGML_CPU_REPACK_MMID_BATCHED=0 in the environment makes off the initial value. -The batched path: thread 0 lists the experts that got rows; the work is (expert, column chunk) -units, about 16 per thread, taken through ggml_threadpool_chunk_add(). In a unit, every group of up -to 16 of the expert's tokens (in fours) is gathered from src1, quantized interleaved with -ggml_quantize_mat_t() -- the quantizer forward_mul_mat() uses for its gemm rows --, multiplied with -gemm() into a per-thread tile and scattered back to the token rows; fewer than four leftover tokens -take the gemv path from the rows quantized before. work_size() reserves the per-op state, the -expert list and the per-thread scratch (16 float rows, their quantized form and a 16-row output -tile; ~420 KB per thread for gpt-oss-20b) whatever the switch says, since it may flip before the -graph runs. +The batched path: thread 0 lists the experts that got rows and where each one's groups of four +tokens start in a shared buffer. Every such group is gathered from src1 and quantized interleaved +with ggml_quantize_mat_t() -- the quantizer forward_mul_mat() uses for its gemm rows -- once, the +groups split statically over the threads, then a barrier (both skipped when no expert got four +tokens, as in token generation). The work is then (expert, column chunk) units, about 16 per +thread, taken through ggml_threadpool_chunk_add(); a unit multiplies up to 16 of the expert's +shared rows at a time with gemm() into a per-thread tile and scatters them back to the token rows, +and fewer than four leftover tokens take the gemv path from the rows quantized before. +work_size() reserves the per-op state, the shared rows (one per routed token slot; 6.3 MB for +gpt-oss-20b at a 512-token batch) and the per-thread scratch (four float rows and a 16-row output +tile) whatever the switch says, since it may flip before the graph runs. + +Quantizing once matters where the activation quantizer is slow: ARM has no SIMD form of +ggml_quantize_mat_q8_K_4x4/4x8 (x86 has one for 4x8), and a first version of this patch, which +quantized an expert's rows again in every column chunk, made Q4_K prompt processing on a +Neoverse-N1 18 % slower at 32 threads (perf: 21 % of the time in the inlined generic quantizer). Measured in java-llama.cpp (docs/benchmarks/moe-mxfp4-cpu-alderlake.md): a random-weight GGUF in gpt-oss-20b's exact shape and ggml-org's MXFP4 layout, the alderlake CPU module (AVX2 + AVX-VNNI) @@ -564,6 +571,17 @@ staged alone, 4 threads: prompt processing +46 % (43.9 -> 64.1 t/s), token gener fell from 13.9 % to 1.7 % of all thread time. llama-perplexity with --kl-divergence against the unchanged code: mean KLD 0.000000, max 0.000002, same top token 100 %. +With real weights (b11901, llama-bench pp512, switch off -> on, t/s; +docs/benchmarks/moe-repack-real-weights.md): a Ryzen 7 5800H (Zen 3, haswell module), 8 threads, +gpt-oss-20b MXFP4 41.7 -> 59.6, Qwen3-30B-A3B Q4_K_M 58.3 -> 84.7, OLMoE-1B-7B Q4_0 +138.8 -> 190.0; a Neoverse-N1 (dotprod, armv8.2_2 module), 32 threads, 177.0 -> 278.7, +190.7 -> 215.7 and 705.9 -> 990.0; an IBM z16 (z16 module, big-endian OLMoE Q4_0), 4 threads, +45.8 -> 76.2. Token generation is unchanged within noise. KLD against the unchanged code on +wikitext-2: 0.000000 for MXFP4 and Q4_0 (same top token 100 %, 99.98 % on the z16); 0.0075 for +Q4_K, where the gemm and gemv kernels round +differently and a router flips experts on near ties -- upstream's repack path is as far from +--no-repack (0.0080 on the N1) as the batched one (0.0085). + The switch is this project's addition, for A/B measurements on a user's own machine (LlamaModel.setCpuMoeBatched, jllama_set_cpu_moe_batched, the environment variable); an upstream submission would offer the batched path alone. diff --git a/llama/patches/0022-ggml-cpu-repack-batch-the-tokens-of-an-expert.patch b/llama/patches/0022-ggml-cpu-repack-batch-the-tokens-of-an-expert.patch index 3da1f2a27..240340797 100644 --- a/llama/patches/0022-ggml-cpu-repack-batch-the-tokens-of-an-expert.patch +++ b/llama/patches/0022-ggml-cpu-repack-batch-the-tokens-of-an-expert.patch @@ -21,15 +21,22 @@ compute cannot send the threads down different paths: (ggml-cpu.h, and in the CPU backend's get_proc_address table for a GGML_BACKEND_DL build). On by default; GGML_CPU_REPACK_MMID_BATCHED=0 in the environment makes off the initial value. -The batched path: thread 0 lists the experts that got rows; the work is (expert, column chunk) -units, about 16 per thread, taken through ggml_threadpool_chunk_add(). In a unit, every group of up -to 16 of the expert's tokens (in fours) is gathered from src1, quantized interleaved with -ggml_quantize_mat_t() -- the quantizer forward_mul_mat() uses for its gemm rows --, multiplied with -gemm() into a per-thread tile and scattered back to the token rows; fewer than four leftover tokens -take the gemv path from the rows quantized before. work_size() reserves the per-op state, the -expert list and the per-thread scratch (16 float rows, their quantized form and a 16-row output -tile; ~420 KB per thread for gpt-oss-20b) whatever the switch says, since it may flip before the -graph runs. +The batched path: thread 0 lists the experts that got rows and where each one's groups of four +tokens start in a shared buffer. Every such group is gathered from src1 and quantized interleaved +with ggml_quantize_mat_t() -- the quantizer forward_mul_mat() uses for its gemm rows -- once, the +groups split statically over the threads, then a barrier (both skipped when no expert got four +tokens, as in token generation). The work is then (expert, column chunk) units, about 16 per +thread, taken through ggml_threadpool_chunk_add(); a unit multiplies up to 16 of the expert's +shared rows at a time with gemm() into a per-thread tile and scatters them back to the token rows, +and fewer than four leftover tokens take the gemv path from the rows quantized before. +work_size() reserves the per-op state, the shared rows (one per routed token slot; 6.3 MB for +gpt-oss-20b at a 512-token batch) and the per-thread scratch (four float rows and a 16-row output +tile) whatever the switch says, since it may flip before the graph runs. + +Quantizing once matters where the activation quantizer is slow: ARM has no SIMD form of +ggml_quantize_mat_q8_K_4x4/4x8 (x86 has one for 4x8), and a first version of this patch, which +quantized an expert's rows again in every column chunk, made Q4_K prompt processing on a +Neoverse-N1 18 % slower at 32 threads (perf: 21 % of the time in the inlined generic quantizer). Measured in java-llama.cpp (docs/benchmarks/moe-mxfp4-cpu-alderlake.md): a random-weight GGUF in gpt-oss-20b's exact shape and ggml-org's MXFP4 layout, the alderlake CPU module (AVX2 + AVX-VNNI) @@ -38,6 +45,17 @@ staged alone, 4 threads: prompt processing +46 % (43.9 -> 64.1 t/s), token gener fell from 13.9 % to 1.7 % of all thread time. llama-perplexity with --kl-divergence against the unchanged code: mean KLD 0.000000, max 0.000002, same top token 100 %. +With real weights (b11901, llama-bench pp512, switch off -> on, t/s; +docs/benchmarks/moe-repack-real-weights.md): a Ryzen 7 5800H (Zen 3, haswell module), 8 threads, +gpt-oss-20b MXFP4 41.7 -> 59.6, Qwen3-30B-A3B Q4_K_M 58.3 -> 84.7, OLMoE-1B-7B Q4_0 +138.8 -> 190.0; a Neoverse-N1 (dotprod, armv8.2_2 module), 32 threads, 177.0 -> 278.7, +190.7 -> 215.7 and 705.9 -> 990.0; an IBM z16 (z16 module, big-endian OLMoE Q4_0), 4 threads, +45.8 -> 76.2. Token generation is unchanged within noise. KLD against the unchanged code on +wikitext-2: 0.000000 for MXFP4 and Q4_0 (same top token 100 %, 99.98 % on the z16); 0.0075 for +Q4_K, where the gemm and gemv kernels round +differently and a router flips experts on near ties -- upstream's repack path is as far from +--no-repack (0.0080 on the N1) as the batched one (0.0085). + The switch is this project's addition, for A/B measurements on a user's own machine (LlamaModel.setCpuMoeBatched, jllama_set_cpu_moe_batched, the environment variable); an upstream submission would offer the batched path alone. @@ -88,7 +106,7 @@ index 81ff5d7..05738c6 100644 // threadpool - TODO: move to ggml-base if (strcmp(name, "ggml_threadpool_new") == 0) { diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp -index 9282ca2..bdeb684 100644 +index 9282ca2..9e9e86a 100644 --- a/ggml/src/ggml-cpu/repack.cpp +++ b/ggml/src/ggml-cpu/repack.cpp @@ -15,9 +15,29 @@ @@ -121,17 +139,22 @@ index 9282ca2..bdeb684 100644 #if defined(__GNUC__) #pragma GCC diagnostic ignored "-Woverlength-strings" #endif -@@ -4554,7 +4574,16 @@ class tensor_traits_base : public ggml::cpu::tensor_traits { +@@ -4554,7 +4574,21 @@ class tensor_traits_base : public ggml::cpu::tensor_traits { template class tensor_traits : public tensor_traits_base { - bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override { -+ // MUL_MAT_ID, batched path: token rows per gemm call, and the scratch each thread needs for one ++ // MUL_MAT_ID, batched path: token rows per gemm call; the rows every expert's gemm reads, quantized ++ // once for all its column chunks (at most one per routed token slot); and each thread's scratch, ++ // four gathered float rows and one gemm's output tile + static constexpr int64_t MMID_GEMM_ROWS = 16; + ++ static size_t mmid_shared_rows(int64_t ne10, int64_t n_rows) { ++ return GGML_PAD(n_rows * ggml_row_size(PARAM_TYPE, ne10), 64); ++ } ++ + static size_t mmid_thread_scratch(int64_t ne10, int64_t ne01) { -+ return GGML_PAD(MMID_GEMM_ROWS * ne10 * sizeof(float), 64) + -+ GGML_PAD(MMID_GEMM_ROWS * ggml_row_size(PARAM_TYPE, ne10), 64) + ++ return GGML_PAD(4 * ne10 * sizeof(float), 64) + + GGML_PAD(MMID_GEMM_ROWS * ne01 * sizeof(float), 64); + } + @@ -139,29 +162,31 @@ index 9282ca2..bdeb684 100644 // not realy a GGML_TYPE_Q8_0 but same size. switch (op->op) { case GGML_OP_MUL_MAT: -@@ -4574,6 +4603,11 @@ template src[1]->ne[0], op->src[2]->ne[0]*ne12); + size += (size_t) n_threads * mmid_thread_scratch(op->src[1]->ne[0], op->src[0]->ne[1]); + return true; } default: -@@ -4845,6 +4879,9 @@ template 0) { -+ mmid_state[2 + n_active++] = cur_a; ++ mmid_state[2 + n_as + n_active] = gemm_rows; ++ mmid_state[2 + n_active++] = cur_a; ++ gemm_rows += matrix_row_counts[cur_a] & ~(int64_t) 3; + } + } -+ mmid_state[1] = n_active; ++ mmid_state[1] = n_active; ++ mmid_state[2 + 2 * n_as] = gemm_rows; + ggml_threadpool_chunk_set(params->threadpool, nth); + } } @@ -190,6 +219,8 @@ index 9282ca2..bdeb684 100644 + // routed to the same expert goes through one gemm instead of a gemv per token. + const int64_t n_active = mmid_state[1]; + const int64_t * active = mmid_state + 2; ++ const int64_t * gemm_row0 = mmid_state + 2 + n_as; ++ const int64_t gemm_rows = mmid_state[2 + 2 * n_as]; + const int64_t n_col_blocks = ne01 / NB_COLS; + if (n_active == 0) { + return; @@ -198,15 +229,39 @@ index 9282ca2..bdeb684 100644 + const int64_t cols_per = ((n_col_blocks + nck - 1) / nck) * NB_COLS; + const int64_t n_units = n_active * nck; + -+ char * scratch = (char *) (((uintptr_t) (mmid_state + 2 + n_as) + 63) & ~(uintptr_t) 63) + -+ ith * mmid_thread_scratch(ne10, ne01); ++ char * shared = (char *) (((uintptr_t) (mmid_state + 3 + 2 * n_as) + 63) & ~(uintptr_t) 63); ++ char * scratch = shared + mmid_shared_rows(ne10, n_ids * ne12) + ith * mmid_thread_scratch(ne10, ne01); + float * fbuf = (float *) scratch; -+ char * qbuf = scratch + GGML_PAD(MMID_GEMM_ROWS * ne10 * sizeof(float), 64); -+ float * obuf = (float *) (qbuf + GGML_PAD(MMID_GEMM_ROWS * nbw1, 64)); ++ float * obuf = (float *) (scratch + GGML_PAD(4 * ne10 * sizeof(float), 64)); + GGML_ASSERT((char *) (obuf + MMID_GEMM_ROWS * ne01) <= (char *) params->wdata + params->wsize); + ++ // every group of four rows routed to one expert is gathered and quantized interleaved (as ++ // forward_mul_mat quantizes its gemm rows) once, here, not once per column chunk: an expert's ++ // rows lie together in the shared rows, so a gemm reads up to four groups in a row. Equal ++ // pieces of work, split statically; skipped with its barrier when no expert got four rows ++ // (token generation) ++ if (gemm_rows > 0) { ++ int64_t group = 0; ++ for (int64_t k = 0; k < n_active; k++) { ++ const int64_t cur_a = active[k]; ++ for (int64_t ir1 = 0; ir1 + 4 <= matrix_row_counts[cur_a]; ir1 += 4, group++) { ++ if (group % nth != ith) { ++ continue; ++ } ++ for (int64_t r = 0; r < 4; r++) { ++ const struct mmid_row_mapping rm = MMID_MATRIX_ROW(cur_a, ir1 + r); ++ memcpy(fbuf + r*ne10, (const char *) src1->data + (rm.i1 % ne11)*nb11 + rm.i2*nb12, ++ ne10*sizeof(float)); ++ } ++ ggml_quantize_mat_t(fbuf, shared + (gemm_row0[k] + ir1)*nbw1, 4, ne10); ++ } ++ } ++ ggml_barrier(params->threadpool); ++ } ++ + for (int64_t unit = ith; unit < n_units; unit = ggml_threadpool_chunk_add(params->threadpool, 1)) { -+ const int64_t cur_a = active[unit / nck]; ++ const int64_t k = unit / nck; ++ const int64_t cur_a = active[k]; + const int64_t c0 = (unit % nck) * cols_per; + const int64_t c1 = MIN(c0 + cols_per, ne01); + if (c0 >= c1) { @@ -216,20 +271,12 @@ index 9282ca2..bdeb684 100644 + const int64_t cne1 = matrix_row_counts[cur_a]; + const char * src0_cur = (const char *) src0->data + cur_a*nb02 + c0*nb01; + -+ // groups of four or more rows: gather the float rows, quantize them interleaved as the -+ // gemm of forward_mul_mat does, multiply, and scatter the rows back ++ // groups of four or more rows: multiply the shared rows, and scatter the result back + int64_t ir1 = 0; + while (cne1 - ir1 >= 4) { + const int64_t nr = MIN(MMID_GEMM_ROWS, (cne1 - ir1) & ~(int64_t) 3); -+ for (int64_t r = 0; r < nr; r++) { -+ const struct mmid_row_mapping rm = MMID_MATRIX_ROW(cur_a, ir1 + r); -+ memcpy(fbuf + r*ne10, (const char *) src1->data + (rm.i1 % ne11)*nb11 + rm.i2*nb12, -+ ne10*sizeof(float)); -+ } -+ for (int64_t r = 0; r < nr; r += 4) { -+ ggml_quantize_mat_t(fbuf + r*ne10, qbuf + r*nbw1, 4, ne10); -+ } -+ gemm(ne00, obuf, nc, src0_cur, qbuf, nr, nc); ++ gemm(ne00, obuf, nc, src0_cur, ++ shared + (gemm_row0[k] + ir1)*nbw1, nr, nc); + for (int64_t r = 0; r < nr; r++) { + const struct mmid_row_mapping rm = MMID_MATRIX_ROW(cur_a, ir1 + r); + memcpy((float *) ((char *) dst->data + rm.i1*nb1 + rm.i2*nb2) + c0, obuf + r*nc,