feat(gguf): serve GGUF checkpoints on vendored ggml-cuda kernels

[jason-fxz]

Part of #671 (GGUF experts, GPU).
Related: #34, #549, #658

Summary

GGUF checkpoints now load through the quant layer like the other formats. A GGUF quant kind reads each tensor's ggml type from the file and runs llama.cpp's CUDA kernels, vendored and bound through tvm-ffi. Qwen3.8-Flash-Next (qwen4exp) is new; Gemma4 GGUF moves onto this path and the old Q4_0-only path is deleted.

Models	Gemma-4 26B-A4B (gemma4), Qwen3.8-Flash-Next (qwen4exp); text only
Tensor types	quantized: Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q2_K to Q6_K, IQ1_S, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ4_NL, IQ4_XS, MXFP4, NVFP4, Q1_0, Q2_0; float: F32, F16, BF16
--moe-strategy	fused, offload; hybrid / cpu for Q4_0 or BF16 experts
TP	1

Vendored ggml

What is copied, from llama.cpp 8c2134a8:

• kernel/csrc/ggml/: the ggml public headers, two ggml-base internal headers, and from ggml-cuda the MMQ / MMVQ / MoE-id / quantize / convert / get_rows kernels with the 22 MMQ template instances.
• third_party/gguf/: gguf-py's gguf_reader, constants, quants and lazy. This drops the gguf PyPI dependency; 0.19.0 lacks type 42 (Q2_0), which Flash-Next uses.

How it is maintained: each tree is upstream plus one freetoken.patch, and its VERSION records the pin, lists the changes and gives the update steps (copy the same files from a newer llama.cpp, apply the patch, regenerate it). 68 of the 138 changed files (32.6k of the 35.7k added lines) are these copies, so reviewing the two patches covers them.

What the ggml patch changes:

• Activations may be fp32, bf16 or fp16; they are quantized to q8_1 directly, and results come out in the activation dtype with fp32 accumulation. MMQ is templated on the output dtype and built for fp32 and bf16; fp16 results go through fp32 unless GGML_MMQ_F16_DST is defined.
• MMVQ can reuse one q8_1 quantization of the input for several weight segments.
• MoE MMQ runs a flat grid over only the (expert, column tile) pairs that hold routes, so it needs no host sync.
• The native fp4 MMA (sm_120) is off unless GGML_CUDA_FP4_MMA is defined; it needs an sm_120a/sm_120f build, and plain sm_120 failed to compile it.

The gguf-py patch trims __init__ to the four modules and makes the reader import its constants relatively.

Our own code next to it: freetoken/ggml_shim.cu (the ggml-base runtime the kernels need) and freetoken/ggml_ops.cu (the tvm-ffi ops: mul_mat, mul_mat_id, fused GLU, MoE weighted sum, get_rows, dequantize). kernel/ggml.py builds it with --compress-mode=size, and the kernel cache includes it.

Changes

• Quant layer: QuantKind.GGUF with Linear, MoE and embedding methods. Fused Linears run one GEMM per run of equal types; MoE uses MMVQ for small batches and MMQ above, with intermediates in the activation dtype (llama.cpp keeps fp32). New LayerKind.EMBEDDING; a tied LM head takes its logits from the embedding method, and quant kinds without an embedding method fall back to unquantized.
• Models: the GGUF reader handles split files and builds config and tokenizer from metadata. qwen4_exp/gguf.py is new, including the PLE table read in place from the GGUF (quantized rows dequantized on the GPU, float rows cast). gemma4/gguf.py is rewritten.
• MoE and engine: banks take each layer's ggml types, layers of different types get separate slot pools, and GGUF experts can be resident (fused).
• Removed: the old Q4_0 path (csrc/gguf/, kernel/gguf.py, layers/gguf.py, models/gguf/dequant.py, moe/fused_q4_0.py).

Validation

tok/s on H100 80GB (hybrid CPU: Xeon Platinum 8480+, 26 vCPUs). AIME25 problem 1, greedy, 256 tokens, one run each. The timed request is the second of two identical requests and reuses the cached prompt (Flash-Next prefills 34 of 98 prompt tokens, Gemma 1 of 63), so the numbers are close to decode speed.

Checkpoint	fused	offload	hybrid
Qwen3.8-Flash-Next GSQ-RCO IQ3_XXS (ISTA-DASLab)	116.0	54.6	-
Gemma-4 26B-A4B QAT Q4_0 (google)	230.1	87.0	110.4
Gemma-4 26B-A4B BF16 (ggml-org)	175.9	49.9*	66.5*
Gemma-4 26B-A4B UD-Q8_K_XL (unsloth, Q8_0 + one BF16 layer)	200.4	70.8	-

offload and hybrid cache 3072 of 24576 experts (12.5%) for Flash-Next and 1024 of 3840 (26.7%) for Gemma. Measured at 371a4c2; the five later commits compress the fatbin, fix reading float PLE tables, pad the offload slot cache for MMQ tail reads, and, from @zihaomu, add HIP build support to the ggml ops (not run on AMD here) and make dequantize copy same-type floats and reject types without a ggml converter.

* BF16 offload was measured before this branch was rebased onto b416f85, BF16 hybrid at 128d0e2.

• fused and offload produce identical tokens. Run to EOS, Flash-Next and Gemma Q4_0 both answer 70.
• Op check against gguf-py dequantization (fp32 / bf16 / fp16): 564 checks pass.
• Builds for the kernel-cache archs (sm_80, 86, 89, 90, 100, 120). Unit tests: 2103 passed (one unrelated remote-checkpoint test deselected).

Not covered

• Tensor parallelism.
• Converting a GGUF checkpoint to FTW.
• NVFP4 token embedding tables: the vendored ggml-cuda get_rows has no NVFP4 case and aborts. Needs a follow-up.
• NVFP4 per-tensor scales: the .scale tensors that llama.cpp's converter writes for weight_scale_2 are not read. Needs a follow-up.
• Q4_0 hybrid decode with more than 8 tokens in a step: CPU routes are parked on slot 0 with weight 0 because the GGUF MoE kernel cannot skip inactive slots, so expert ids repeat within a token and MMQ hits an illegal memory access. The hybrid numbers above are bs=1. Needs a follow-up.
• CPU / hybrid experts other than Q4_0 and BF16: follow-up PR.
• Vision towers (VLM): follow-up PR.
• The native fp4 MMA.

[jason-fxz]

GGUF experts in the offload slot cache have no MMQ tail padding

ggml's MMQ kernels read past the last row of the weight matrix they get, up to ggml_row_size(type, 512 - K % 512) bytes. llama.cpp allocates that tail and zeroes it (ggml_backend_cuda_buffer_type_get_alloc_size and ggml_backend_cuda_buffer_init_tensor), and kernel/ggml.py checks for it in _check_mmq_tail before every MMQ mul_mat_id.

The resident (fused) banks have it: _device_bank appends a zero row. The offload slot cache does not: each bank holds exactly the pools' slots, packed back to back, and the two prefill staging buffers sit at the start of the bank.

Reproduced at 128d0e2 with --moe-strategy cpu on gemma-4-26B-A4B-it-qat-q4_0, which fails at the first prefill:

ValueError: ggml type 2 K=2816 weight needs 144 bytes of zero padding after its last row for MMQ

In cpu mode the cache is exactly 2 x num_experts slots (256 here), so staging buffer 1 ends at the end of the bank. Offload and hybrid with 1024 slots run fine on the same checkpoint.

Proposed fix, for GGUF layouts with K % 512 != 0 only (the tail is 0 for everything else):

• After each pool's region in each bank, reserve mmq_tail_bytes(type, K) bytes rounded up to 16, zeroed at allocation and never handed out as a slot.
• Grow each staging buffer by the largest tail, so the slots it invalidates during prefill cover it, and zero the tail after each layer is staged: the staging buffers alias pool slots and can hold bytes of a wider layout.

[zihaomu]

Clarification of the exact commit sets from the single-GPU AMD integration branch: zihaomu/FreeToken:feat/pr690-pr691-rocm-compat, head 8a693819285d6eea9575cda1540b92038aca0418. The branch combines the original #690 and #691 heads plus the seven commits listed below; it has not yet been rebased onto current upstream main.

Base prerequisite: rebase #690 onto current main, including #691, upstream commit 0bca88f70c4909624f7f3c3e3a4fe6b32d183d56. This provides the ROCm/Triton 3.8 MoE alignment workaround. Keep that version-gated staged path when resolving overlaps. The runtime foundation #132 is already merged. Remaining runtime work in #587 and CPU/hybrid graph work in #667 should be rebased/reconciled with #670 and #691; merging either entire PR is not an established prerequisite for the GPU-only GGML path tested here.

A. ROCm-specific enablement and correctness commits

For the tested R9700 (gfx1201), Torch 2.14 / Triton 3.8 stack, the concrete production integration set is #691 via rebase + 75f9c0f + f1826a5 + 034e52a, with the path conditions below. Carry 8a69381 as regression coverage.

Full commit SHA	Purpose	When it is required
75f9c0fb8ea1561fc7db9909f3786c38d0d0682c	fix(gguf): support hip builds and runtime devices	Required to build/run this new GGML backend on HIP: HIP flags, ROCm DLPack devices/allocations, default-stream handling and AMD device identification.
f1826a51040ecbf2f0e88de73cff9f1f4df58d1d	fix(qsa): keep rocm attention tiles within lds limits	Required for the tested Qwen QSA D=256 prefill on the R9700: the original tile requested 65,792 bytes against 65,536 bytes of LDS. Reconcile the overlapping QSA change in #587.
034e52a0513590b4dbfb746746924351153a7332	fix(rocm): round bf16 moe products on affected triton builds	Required for correct BF16 product rounding in dedicated MoE decode on affected gfx11*/gfx120* compiler builds. A cached compile-only probe enables it only when the bad lowering is present; fixed compilers and CUDA skip it.
8a693819285d6eea9575cda1540b92038aca0418	test(rocm): preserve cuda coverage in backend regressions	Tests only, not a production runtime prerequisite. Preserve NVIDIA FlashInfer/CUDA 13 coverage while adding ROCm-compatible checks, Gemma router ties, and alignment/graph replay coverage.

The BF16 workaround is separate from #691: it addresses Triton #11217 / #11227, where BF16 multiplication is lowered to truncating v_dot2_bf16_bf16. It does not switch resident/grouped dispatch or change reduction precision. On 1,048,576 pairs, CPU RTNE agreement improves from 52.60% to 100%. Its synthetic Gemma-shape BF16 MoE decode cost was 0.1783 to 0.1944 ms (~9%); this is not an end-to-end slowdown measurement.

B. General GGUF/GGML correctness fixes, reviewed independently of ROCm

These address format, storage and tokenizer behavior shared by CUDA and ROCm. The two defensive candidates left for design discussion are 6cc138e and b6c7e66. There is also the separate conversion fix 90f9833, listed here so all seven branch commits are accounted for.

Full commit SHA	Purpose	Evidence / review scope
90f9833ac3feea9d59d99c8ecd2a67b1d4144a33	fix(gguf): handle nvfp4 lookups and identity conversions	General GGML conversion correctness: NVFP4 row lookup fallback and same-dtype float copies instead of calling a null converter. Relevant when these formats/operations are used, independently of GPU vendor.
6cc138eb0635fcf6da07c134046a341e1bd5b27f	fix(gguf): pad offload cache storage for mmq tails	Defensive candidate for discussion. With 128 experts / 256 slots, staging buffer 1 reaches the bank end; Q4_0/K=2816 needs 144 extra bytes and fails the existing Python guard. The candidate adds a zeroed 2 KiB storage tail without adding slots. Reconcile it with the existing per-pool/staging-tail reservation and zeroing proposal, especially for aliased staging and mixed layouts.
b6c7e6613015461ada064dbabe2676a291b33994	fix(gguf): preserve control and user-defined token ids	Defensive candidate for discussion. Register metadata-defined atomic tokens. The tested Qwen <think> is USER_DEFINED, ID 248068; original conversion splits it into [13314, 741, 29], while registration restores [248068] and the same prompt changes from 34 to 30 tokens. Full vocabulary IDs and existing added-token settings remain unchanged on the tested Qwen/Gemma checkpoints. The directly demonstrated need is USER_DEFINED; review the breadth of the additional CONTROL registration separately.

Putting a fix in B does not make its triggered failure optional: the tested Qwen tokenizer and small-cache MMQ cases still need a solution. The separation is that these are general GGUF bugs with independently discussable implementations, not ROCm-only dependencies.

Latest targeted validation: one R9700, TP=1, Torch 2.14.0+rocm10.1.0, HIP 7.16.26385, Triton 3.8.0+git669b31ac.rocm10.1.0; 127 targeted tests passed. Gemma UD-Q8_K_XL offload passed two matching 256-token requests and a full EOS request answering 70. Eight NVIDIA offline compilations (SM80/90/100/120, FP16/BF16, workaround disabled) produced identical kernel machine-code sections. No NVIDIA hardware execution or complete-suite pass is claimed. The history cleanup preserved the full tracked tree used for these results.

[zhangnju]

Reconciling the remaining runtime work in #587 with the ROCm integration set here.

QSA LDS clamp — overlaps f1826a5. #587's qsa/attend.py change solves the same RDNA LDS overflow as f1826a5 (the D=256 / N=64 K+V tile exceeding 64 KiB). The two converge on the tested case — both clamp block_n 64→32 for D≥256 bf16. One difference worth folding in: #587 queries the device's actual shared_memory_per_block and shrinks only when the tile truly does not fit, so a larger-LDS AMD part (e.g. gfx950) keeps the wide tile, whereas a hardcoded head_dim >= 256 -> min(block_n, 32) clamps it unconditionally. Proposal: keep f1826a5's torch-reference test (#587 has none) and adopt the device-query budget so the clamp generalizes past the fixed 64 KiB assumption. Happy to push that onto the integration branch or hand the hunk over, whichever is easier.

gguf fixes — superseded by this PR. #587's other two changes — csrc/gguf/dispatch.h (64-bit __shfl_xor_sync warp mask for the amd_warp_sync_functions static_assert) and gguf.py (stripping the nvcc-only --expt-relaxed-constexpr / -ccbin on ROCm) — are on files this PR deletes. The old csrc/gguf/ tree is replaced by the vendored csrc/ggml/ tree, whose HIP build support is 75f9c0f. So those two are moot under #690.

Net: once #690 + the ROCm integration set land, #587 is fully covered — QSA by f1826a5 (ideally folding in the device-query budget), gguf by 75f9c0f — and I'll close #587.

[jason-fxz]

@zihaomu
Thanks for splitting the integration branch into separate commits. Here is where each one goes.

Taken into this PR (with your authorship):

• 75f9c0f (HIP build) as aeed072. On CUDA the build flags are unchanged and tests/kernels/test_ggml.py passes on an H100; it has not been run on AMD here.
• The op_dequantize part of 90f9833 (same-type copy and null converter checks) as a3558d2.

Not taken:

• 6cc138e: 1d8a5df fixes the same MMQ tail bug, and also covers pool boundaries and staging buffers that overlap slots of another layout.
• b6c7e66: we prefer to fix the tokenizer by moving to transformers 5.18.0, which registers both USER_DEFINED and CONTROL tokens, in a separate PR.
• Rebasing onto main: fix(rocm): reuse staged moe alignment on Triton 3.8 #691 only changes moe/fused.py, which this PR does not touch, and the PR merges cleanly.

For a new ROCm PR from you against main: f1826a5 (reconciled with the QSA change in #587), 034e52a and 8a69381. It does not depend on this PR, except the test_offload.py part of 8a69381, which changes a file this PR also changes.

On hold: the rest of 90f9833 (NVFP4 get_rows fallback and tests). NVFP4 GGUF needs more than the lookup: the .scale tensors that llama.cpp's converter writes for weight_scale_2 are not read either. We would like to check NVFP4 end to end on a real NVFP4 GGUF checkpoint first. Both gaps are listed under "Not covered" in the description.

[jason-fxz]

Four fixes on top of a3558d2, tested on an RTX PRO 6000 Blackwell (sm_120).

8682d4e – quantized PLE projections (qwen4exp)
The loader read ple_key / ple_value as float, and ISTA's Q2_0 file stores ple_key as Q2_0, so loading failed with KeyError: 42. Both projections now go through fused() / matrix() like the other Linears. Across the Flash-Next GGUFs on the Hub, these two tensors are BF16 (ISTA IQ*), Q2_0 (ISTA Q2_0) or Q8_0 (every unsloth quant, AtomicChat).
Tested: ISTA Q2_0, fused: AIME#1 = 70, 147.4 tok/s at bs=1.

0918f1e – float parts in a fused GGUF Linear
A fused module whose parts mix quantized and float tensors raised NotImplementedError. unsloth's UD-* files hit this: hc_*_down is Q8_0 while hc_*_inject is F32. A float run of parts is now its own bf16 segment, written in place with torch.mm(out=column slice), and the quantized segments run on ggml as before. fused() emits one tensor per segment and puts pad_to zero rows on the last segment, which it used to drop for quantized parts. All-float and all-quantized modules behave as before.

dfbe98f – quantized hyper-connection and GDN ba weights (qwen4exp)
hc_*_{down,inject,up}, output_hc_* and ssm_alpha/beta were read as float. unsloth and AtomicChat store the hc matrices as Q8_0 (and ssm_alpha/beta as Q8_0 in unsloth Q8_0 and AtomicChat). They now load through matrix() / fused(). LinearReplicated passes output_sizes through, and the merged hc Linear declares (lowrank, hc_count + pad) so the GGUF method can split its segments.
Tested:

• Loader dry run against the meta-built model: unsloth UD-IQ1_M (923/923 params) and ISTA Q2_0 (899/899) match in key, shape and dtype.
• unsloth UD-Q2_K_XL serves under fused (about 140 tok/s at bs=1).
• Not run end to end: unsloth Q8_0 and AtomicChat (Q8_0 ssm_alpha/beta).

64ef88f – NaN after prefill with offload
1d8a5df pads the end of each pool and staging buffer. But MMQ also reads past every slot's last row into the next slot's first bytes. After prefill, the invalidated slots under the staging buffers still hold staged rows that are shifted by the pad (one quant) or belong to another quant. A routed slot next to such a slot then reads a garbage block scale, and NaN × 0 is still NaN.
Repro: Gemma Q4_0, offload --moe-cache-rate 0.267, 96 staggered requests at 16 concurrency: 23/96 outputs contain <pad>.
Fix: at the last prefill layer, zero the byte range of every staging buffer used in this prefill, in every bank. This replaces the pool-pad and slot-0 row lists and the zero_inactive_route_slots flag. The cost is one memset of up to two staging buffers per bank per prefill (about 1 GB, under 1 ms).
Tested:

• Gemma Q4_0 offload: 0/96 with <pad>.
• Gemma UD-Q4_K_XL offload: AIME#1 prompt and 24 staggered requests clean.
• HF Gemma-4-26B-A4B-NVFP4 offload: 48 staggered requests at 16 concurrency clean.
• tests/moe: 152 passed.

Still open, not in these commits

• The GGUF tokenizer imports transformers.integrations.ggml.convert_gguf_tokenizer, which 5.18 moved to transformers.integrations.gguf; on the pinned 5.16, USER_DEFINED tokens such as <think> are split.
• hybrid parks CPU routes on slot 0, so ids repeat within a token and MMQ hits an illegal memory access above 8 tokens.
• MMQ / MMVQ use 32-bit block offsets, so a bank over 2^31 blocks faults.
• IQ1_M experts have no MMQ kernel above 8 tokens.
• Qwen3.8 GGUF decode at 16 concurrency hangs in graph replay (separate PR, problem with PLE sync-wait).