Skip to content

feat(kernel): add an iq1_m mmq kernel for moe steps above 8 tokens - #835

Merged
jason-fxz merged 1 commit into
FlashML-org:mainfrom
ischencheng:feat/gguf-iq1m-mmq
Oct 29, 2026
Merged

jason-fxz merged 1 commit into
FlashML-org:mainfrom
ischencheng:feat/gguf-iq1m-mmq

Conversation

@ischencheng

Copy link
Copy Markdown
Contributor

Fixes the IQ1_M known issue in #671: IQ1_M experts have no MMQ kernel, so any MoE step above the MMVQ batch (8 tokens on Ada and Blackwell) fails with no ggml MoE kernel for iq1_m with T=..., and the dense path fails the same way. So on fused and offload every prompt longer than 8 tokens fails.

IQ1_M is not rare in the files people publish. I read the headers of 125 IQ1 files in 68 HF repos: 83 are MoE files of an architecture FreeToken serves, and 59 of those have IQ1_M experts, including unsloth's GLM-5.3-Flash UD-IQ1_M, MiniMax-M3.1-Flash UD-IQ1_M, GLM-4.5/4.6/4.7, MiniMax-M2.x/M3, Qwen3.6-35B-A3B-MTP and Qwen3-235B. 21 of the 35 files named IQ1_S also have some IQ1_M expert layers. llama.cpp master has no IQ1_M MMQ either; its mul_mat_id falls back to one ggml_cuda_mul_mat per expert there (dequantize + cuBLAS), which FreeToken's MoE path does not have.

IQ1_S's MMQ loader cannot be reused: it folds the delta into a per-32 dm pair and the q8_1 block sums, while IQ1_M's delta changes every 8 values and its scale every 16. But 8 * (g + delta) is always an integer (-9, -7, -1, 1, 7, 9), so the new loader writes those as int8 with d * (2s + 1) / 8 per 16 values. That is exactly IQ2_XS's tile (GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K) and dot product (vec_dot_q8_0_16_q8_1), with no rounding on the weight side.

  • mmq-load-tiles.cuh: ggml_cuda_mmq_load_tiles_iq1_m, written like the IQ2_XS loader, for both the MMA and the DP4A layout.
  • mmq.cuh: IQ1_M in the DP4A and MMA dispatch tables, D4 activation layout, Q8_0_16 DP4A tile size. ggml_cuda_mmq_get_config gives IQ1_M the IQ2_XS configs (same tile layout and size), so the per-arch config files are unchanged.
  • mmq.cu: IQ1_M in the type switch and in ggml_cuda_should_use_mmq.
  • template-instances/mmq-instance-iq1_m.cu: the kernel instance, picked up by kernel/ggml.py's glob.
  • freetoken.patch is regenerated against llama.cpp 8c2134a8 (the same git diff --no-index gives today's patch byte for byte on main, and the new patch applied to the 8c2134a8 files gives the vendored tree byte for byte), and VERSION lists the change.

test_iq1_m_matrix_and_moe_routes_match_cpu runs mul_mat and mul_mat_id on random IQ1_M blocks (every grid index, shift bit and 3-bit scale) at 1, 17 and 300 tokens, against gguf-py's dequantization. gguf-py has no IQ1_M quantizer, so _iq1_m_weights builds the blocks directly, like _nvfp4_weights. The fp32 cases take random fp16 super-scales. The bf16 / fp16 cases take m / 1024 ones, so with _activation every partial sum stays exact in fp32 and the rounded result matches the reference.

Tested on A10 (sm_86) and L40S (sm_89), driver 610.57.04, CUDA 13.0, torch 2.14.1+cu130, Triton 3.8.0, on top of 9221fd7.

pytest tests/kernels/test_ggml.py::test_iq1_m_matrix_and_moe_routes_match_cpu, each case in its own process: on main (only the test applied) the 3 single-token cases pass and the 6 at 17 and 300 tokens fail with no ggml kernel for iq1_m with M=17 / M=300; with this PR all 9 pass on both GPUs. With the (sc[0] >> 12) term dropped from the loader, the fp32 cases at 17 and 300 tokens fail (92% of the elements off), so the super-scale's low nibble is covered.

python scripts/gguf_check/probe_kernels.py on A10: the only row that changes is IQ1_M, whose mm16, mm512, mmid16 and mmid512 go from the runtime check error to ok.

python scripts/gguf_check/check.py --kernels <that table> on header stubs of 16 IQ1 files (31 shards) from unsloth and bartowski:

main this PR
pass 3/16 16/16
first problem of the others 12x experts gate_up IQ1_M: no mmid16, mmid512 kernel, 1x linear IQ1_M: no mm16, mm512 kernel -

The 16 files: GLM-4.5-Air UD-IQ1_M/S, GLM-5.3-Flash UD-IQ1_M/S, DeepSeek-V4-Flash-0731 UD-IQ1_M/S, MiniMax-M2.5 UD-IQ1_M/S, MiniMax-M3.1-Flash UD-IQ1_M, Qwen3-30B-A3B-Instruct-2507 UD-IQ1_M/S, Qwen3-Coder-30B-A3B-Instruct UD-IQ1_M/S, Qwen3.6-35B-A3B UD-IQ1_M (unsloth, plain and MTP) and bartowski's Qwen3.6-35B-A3B IQ1_M. The three that pass on main have no IQ1_M tensor.

End to end on L40S with unsloth/Qwen3.6-35B-A3B-MTP-GGUF, file Qwen3.6-35B-A3B-UD-IQ1_M.gguf (gate/up experts IQ1_M in most layers, IQ2_XXS / Q2_K in the rest):

ft serve --model Qwen3.6-35B-A3B-UD-IQ1_M.gguf --moe-strategy fused --max-running-requests 16 --max-seq-len-override 4096
ft serve --model Qwen3.6-35B-A3B-UD-IQ1_M.gguf --moe-strategy offload --moe-cache-rate 0.25 --max-running-requests 16 --moe-collect-stats --max-seq-len-override 4096

then 4 greedy prompts longer than 8 tokens, streamed one at a time (64 tokens each), then 16 greedy /v1/completions sent together (Question: what is a * b?\nAnswer: a * b =).

fused offload
main does not start: capturing the bs=16 decode graph fails with no ggml MoE kernel for iq1_m with T=16 not run (same kernel)
this PR starts; 197 tok/s at bs=1; 15/16 correct starts; 87 to 114 tok/s at bs=1, the 4 streamed texts identical to fused; 15/16 correct

The one wrong product (66 * 44 -> 2916) is the same on both strategies. No CUDA errors in either log. --moe-strategy cpu refuses this file on main and with this PR alike (experts in several quant formats need the GPU slot cache), so I had no cpu-side output to compare against.

A larger file, on the same L40S (48 GB) with 175 GB of host RAM: unsloth/GLM-5.3-Flash-GGUF, UD-IQ1_M/GLM-5.3-Flash-UD-IQ1_M-00001-of-00003.gguf (91 GB in 3 shards, gate/up experts IQ1_M or IQ2_XXS), --moe-strategy offload --moe-cache-rate 0.2 --max-running-requests 16, the same requests (48 tokens per streamed prompt): it starts, the 4 prompts come back coherent, 16/16 products are correct, and the log has no CUDA errors. On main check.py already reports this file as experts gate_up IQ1_M: no mmid16, mmid512 kernel, so I did not start it there.

pytest tests/kernels/test_ggml.py tests/moe tests/models/test_flint.py tests/models/test_gemma4_gguf_rope.py -m "not slow": main 193 passed / 23 skipped on A10; this PR 207 passed / 18 skipped on A10 and L40S, no failures. The difference is the 9 new cases and 5 test_cpu_moe_fp8_block.py cases that skip without AVX-512 and ran here because those hosts had it; the other skips are the same.

Not tested: AMD (the MFMA/WMMA branch follows IQ2_XS's), Pascal DP4A, and IQ1_M speed tuning (it runs on IQ2_XS's tile configs). hybrid above 8 tokens per step still needs #823, as for the other quantized types.

@jason-fxz

Copy link
Copy Markdown
Collaborator

LGTM. I also ran it on an RTX PRO 6000 (sm_120), which the PR did not cover:

  • tests/kernels/test_ggml.py: 31/31 pass.
  • Accuracy on real expert shapes (rows a multiple of 128 and K=768 with an MMQ tail, 9 to 1024 tokens, fp32 and bf16): NMSE against x @ dequantize(w).T is ~3e-5, the same as IQ2_XS on the same shapes, so it is all q8_1 activation rounding.
  • Speed: IQ1_M MMQ is within 5% of IQ2_XS MMQ on every shape. On MoE it is 1.6-17x faster than the per-expert dequantize + cuBLAS route llama.cpp takes for IQ1_M (Qwen3.6-35B-A3B gate_up, 256 tokens: 0.32 ms vs 4.0 ms).

@jason-fxz
jason-fxz merged commit d2977fb into FlashML-org:main Oct 29, 2026
1 check passed
ischencheng added a commit to ischencheng/FreeToken that referenced this pull request Oct 29, 2026
Resolves freetoken.patch and VERSION after FlashML-org#835: the patch is regenerated against llama.cpp
8c2134a8 with both changes, and VERSION lists both.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants