feat(cpu-moe): serve GGUF experts on the cpu and hybrid backends through ggml-cpu

[jason-fxz]

Part of #671. Builds on #700.

Summary

GGUF experts of every ggml type, quantized or float, now run on the CPU MoE executor (hybrid, and cpu / --moe-cpu-layers where the slot pools allow), not only Q4_0 and BF16. This includes layers whose gate_up and down mix a float and a quantized type (#702). The kernels are llama.cpp's ggml-cpu row kernels, vendored unchanged and built once per x86 variant of llama.cpp's GGML_CPU_ALL_VARIANTS; at startup the executor uses the best variant this CPU runs. The q4_0 CPU format is removed.

How it plugs into the CPU executor

GGUF is one more weight format. The executor (threads, work split, route sum) has nothing GGUF-specific.

Startup:

Engine._init_cpu_moe_executor
  -> GGUFMoEMethod.cpu_spec()             per layer: ("gguf", {gate_up_type, down_type}), or ("bf16", {}) when both are BF16
  -> CpuMoeExecutor(layer_formats=...)    one format instance per distinct spec (from #700)
  -> ggml_cpu.select()                    once per process: imports each variant module, keeps the highest score()
  -> _cpu_moe.set_ggml_cpu(api)           the chosen variant's kernel table
  -> _resolve_gguf_banks                  checks each bank holds rows of row_bytes(type, K)
                                          returns the per-layer gate_up / down bank tables and {gate_up_type, down_type}
  -> _cpu_moe.CpuMoeExecutor(formats, layer_formats)
  -> make_format("gguf", params, FormatArgs{banks, options})
  -> GgufFormat(params, args)             picks vec_dot and the input quantizer for each of the two types from that table

Decode (prefill stays on the GPU):

OffloadMoELayer._decode_routed
  -> cpu layer: CpuMoeExecutor.decode                          all routes on the CPU
     hybrid:    _decode_hybrid -> decode_submit / decode_sync   missed routes on the CPU, added to the GPU result
  -> submit:  GgufFormat::prep_input   bf16 input -> the gate_up type's input (q8_0 / q8_1 / q8_K, or f16 / bf16 / f32)
  -> workers: do_pass1 -> GgufFormat::gate_up -> ggml_vec_dot_<type>* per gate and up row, then the activation
              GgufFormat::prep_inter   intermediate row -> the down type's input
              do_pass2 -> GgufFormat::down    -> ggml_vec_dot_<type>* per output row and route

Rows are read in place from the host banks the GPU copies from. At startup gguf is refused only when the build has no variant module (not x86-64), with an error pointing to --moe-strategy offload. The startup log names the variant, for example isa=zen4(ggml-cpu), and FREETOKEN_GGML_CPU_VARIANT=<name> forces one.

CPU variants

• One Python extension module per variant, freetoken.kernel.ggml_cpu._<variant>: the 14 x86 variants of upstream GGML_CPU_ALL_VARIANTS (x64, sse42, sandybridge, ivybridge, piledriver, haswell, skylakex, cannonlake, cascadelake, icelake, cooperlake, zen4, alderlake, sapphirerapids), or the 9 of them upstream builds with MSVC. Each module holds the kernel files built with that variant's ISA flags and macros, as upstream's CMake sets them, and an entry (variant_module.cpp with upstream's cpu-feats.cpp) built without them.
• select() imports every module and keeps the highest score(), upstream's rule: the score is 0 when the CPU lacks an ISA extension the variant is built for. Importing a module runs only the entry. x64 runs on any x86-64 CPU, so GGUF experts now also run on CPUs without AVX2.
• ggml-quants.c is built once without ISA flags, like upstream's ggml-base, and linked into every module.
• Our own formats in _cpu_moe keep their per-function target attributes and runtime dispatch. _cpu_moe no longer links any ggml code.
• The 14 modules are 376-457 KB each after strip (5.6 MB in all; 25.5 MB with the -g debug info setuptools adds).

Vendored ggml-cpu

From llama.cpp 8c2134a8, the same commit as the ggml-cuda tree of #690. No file is modified; VERSION lists them.

• Copied: only the row-level kernels, ggml-cpu/quants.c and ggml-cpu/arch/x86/quants.c (per-type vec_dot and input quantizers), ggml-cpu/vec.cpp (the f32 / f16 / bf16 vec_dot) and ggml-quants.c (reference quantizers and grids), with their headers, and ggml-cpu/arch/x86/cpu-feats.cpp (the variant score) with ggml-backend-impl.h. The ggml-cpu backend (ggml-cpu.c, ops, threadpool, graph) is not copied; the executor keeps its own threads and calls the kernels row by row.
• Glue, ggml/freetoken/cpu_kernels.c: the few ggml-base symbols the kernels need, the fp32 -> f16 / bf16 / f32 input converters, and a weight type -> vec_dot / input type table whose 28 entries match type_traits_cpu in upstream ggml-cpu.c.
• Update: copy the same files from a newer llama.cpp, update VERSION, add new types to the table, and new upstream variants to GGML_CPU_VARIANTS in setup.py.

Validation

H100 80GB, driver 580.95.05, CUDA 13.0; Xeon Platinum 8480+, 12 executor threads; torch 2.14.1+cu130, gcc 11.4.

python setup.py build_ext --inplace
pytest tests -m "not slow" --ignore=tests/e2e   # 2123 passed, 2 failed (below)

The two failures are not from this PR: one is a local HF cache entry with an unsupported modelopt_mixed config, and test_offload.py::test_graph_capture_reuses_warm_offload_cache_before_capture (FakeGraph has no raw_cuda_graph_exec) also fails on main e8c6bf1.

• Kernel check: CpuMoeExecutor.decode against an fp32 reference built from the dequantized rows, bs 1, 3 and 8, top-8, with the variant select() picks on the test CPU (zen4). Worst max error relative to the reference max:

  rows	CPU	GPU
  Flash-Next, real (IQ2_XS, IQ2_XXS, IQ2_S, IQ3_S, IQ3_XXS gate_up; IQ4_NL, Q2_0 down)	1.7e-2	1.4e-2
  Gemma, real (Q4_0, Q8_0)	1.4e-2	3.7e-2
  Gemma shape, synthetic (Q4_1, Q5_1, Q4_K, Q5_0, Q6_K, MXFP4, Q5_K)	2.4e-2	3.0e-2
  Gemma shape, synthetic float (F16 + Q8_0, BF16 + Q4_0, Q4_K + F16, Q6_K + BF16, F16, F32, BF16)	1.3e-2	2.6e-2

  The GPU column runs the GPU path of the same layers: the ggml-cuda kernels for quantized rows, and the fused MoE Triton kernels for float rows (feat(gguf): run experts whose gate_up and down mix float and quantized types #702).

• The same check with each of the other 12 variants the test CPU runs, forced with FREETOKEN_GGML_CPU_VARIANT (sapphirerapids scores 0: the VM does not expose AMX): all pass. The worst CPU error is 2.4e-2, and 2.6e-2 for x64 and sse42, which run ggml's scalar code.

• Q4_0 decode on the CPU, Gemma 26B-A4B expert shape with real rows from the QAT Q4_0 GGUF, 8 layers x 32 experts, top-8, 12 threads. Median us per layer, two runs:

  	bs 1	bs 4	bs 8
  main, q4_0 format (its own VNNI kernel)	402 / 369	1404 / 1342	2659 / 2605
  haswell (AVX2, what this PR built before)	420 / 411	1473 / 1486	2843 / 2854
  skylakex (AVX-512, no VNNI)	412 / 405	1488 / 1483	2887 / 2886
  alderlake (AVX2 + AVX-VNNI)	392 / 390	1491 / 1403	2719 / 2720
  cascadelake (AVX-512 VNNI)	397 / 395	1408 / 1415	2729 / 2756
  zen4 (picked on this CPU)	395 / 396	1416 / 1420	2753 / 2749

  Q8_0 rows (unsloth Q8_K_XL GGUF) take 505 to 537 us at bs 1 on each of the five variants above.

• A CUDA graph decoding a GGUF layer and a BF16 layer replays identically to eager decode.

• Windows, compile only: clang-cl with the xwin MSVC SDK builds each of the 9 MSVC variants with the MSVC flags in setup.py (plus the -m flags clang needs for VBMI, VNNI and AVX-VNNI, as upstream adds for Clang) and gguf.cpp, and lld-link links each variant with no unresolved symbols. The module entry was built against a stub Python.h, since no Windows Python headers were at hand.

Not covered

• End to end at this head.
• cpu / --moe-cpu-layers on GGUF files whose layers use several ggml types: rejected at startup, since several slot pools do not combine with them.
• A CPU without the extensions of a variant: each variant ran on the test CPU only by forcing it. The module entries, which run at import on any CPU, are built without ISA flags.
• Whether upstream's score order is the fastest choice for every type: in two runs, one of them disturbed by another process, the picked zen4 was 2 to 6% slower than haswell at bs 8 for IQ2_XXS / Q2_0 and Q4_K / Q5_0 rows, about the run-to-run noise there.
• ft bench bw for GGUF (hybrid uses a fetch cap of 1), one kernel call for several tokens of an expert, aarch64, running on Windows, and wheel installs of the variant modules (this was built with build_ext --inplace).

[andy-yang-1]

The design looks good to me: only the row kernels are vendored, and the #700 executor keeps the threads and the work split, so GGUF is just one more format. I found no correctness bug. Before merge I would add tests and state two regressions in the PR.

Checked (H100, Xeon 8480+, torch 2.9.0+cu128):

• All 11 vendored files are byte-identical to llama.cpp 8c2134a8, and the 28 vec_dot / input-type entries in cpu_kernels.c match upstream type_traits_cpu.
• The AVX2 objects have no .init_array, so importing on a CPU without AVX2 is safe. No ggml symbol is exported from _cpu_moe, and the MSVC flags match upstream.
• My own check against an fp32 reference built from the dequantized rows covered 20 type pairs in one executor (Q4_0/1, Q5_0/1, Q8_0, MXFP4, Q2..Q6_K, IQ1_S..IQ4_XS, and F16/BF16/F32 alone or next to a quantized type) at bs 1, 3 and 8. The worst max error relative to the reference max is 2.0e-2.
• The wheel build and an editable install both produce a _cpu_moe that lists gguf. That item can come off "Not covered".
• Full -m "not slow" suite (without e2e): 2081 passed, 45 failed, 2 errors. All 47 also fail on main e8c6bf11, so there are no regressions.

1. No tests. The PR deletes test_cpu_moe_q4_0.py (decode vs dequant + GPU, vs ggml MMVQ, CUDA graph replay) and adds none; the kernel table in the description is not committed. Those tests port to the gguf spec with 3 changed lines, and all 5 cases pass. Please keep them, and add a test that runs one layer per type pair against the dequantized reference.

2. Q4_0 is 6-13% slower than main, and there are no A/B numbers. Setup: Gemma 26B-A4B expert shape, 8 layers x 32 experts (~0.9 GB, well past the LLC), top-8, 12 threads, same banks. Median µs per layer over 2 runs:

bs	main q4_0 (VNNI)	this PR (AVX2)	this PR built with -mavxvnni
1	348-352	374-387	354-359
4	1256-1260	1328-1402	1280-1314
8	2421-2446	2565-2740	2575-2627

So VNNI does help on this CPU. Options: an AVX-VNNI build dispatched at runtime (the cpuid check is already there), or accept the loss and state it in the PR.

3. aarch64 loses GGUF on the CPU. The removed q4_0 format had a scalar fallback, so Q4_0 GGUF ran on cpu/hybrid on Grace (GH200/GB200) and Jetson. FT_GGML_CPU is x86-only, so every GGUF type except BF16 is now refused there. "aarch64" under "Not covered" reads like support that was never there, not support being removed. Vendoring upstream ggml-cpu/arch/arm/quants.c (NEON) would bring it back and make it faster.

Nit: _cpu_moe_executor_viable now returns True for GGUF that uses several ggml types. --moe-cpu-layers auto then picks layers and only later stops with the several-slot-pools error, instead of failing up front.