Skip to content

[CUDA] Add dense and MoE GGUF MMQ kernels for eight I-quant types - #36122

Draft
liuhao-labs wants to merge 3 commits into
sgl-project:mainfrom
liuhao-labs:fix/issue-35019-iq-mmq
Draft

[CUDA] Add dense and MoE GGUF MMQ kernels for eight I-quant types#36122
liuhao-labs wants to merge 3 commits into
sgl-project:mainfrom
liuhao-labs:fix/issue-35019-iq-mmq

Conversation

@liuhao-labs

@liuhao-labs liuhao-labs commented Aug 24, 2026

Copy link
Copy Markdown

Motivation

SGLang supports I-quant GGUF weights for dequantization and MMVQ, but does not
currently include them in dense or batched-MoE MMQ dispatch. Beyond the
small-batch MMVQ boundary, dense I-quant prefill dequantizes the full weight to
BF16, while I-quant MoE prefill remains on the decode-oriented vector kernel.

This change adds the eight I-quant formats supported by the upstream MMQ
implementation. IQ1_M remains on fallback because that implementation does
not provide its MMQ dot product.

Fixes #35019

Modifications

  • Add CUDA dense and batched-MoE MMQ support for IQ1_S, IQ2_XXS,
    IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ4_NL, and IQ4_XS.
  • Apply a conservative M<=16 IQ MMQ gate. Wide matrices use MMVQ through M=8,
    MMQ for M=9..16, and dequantization from M=17. Narrow matrices retain MMVQ
    through M=16 and dequantization from M=17.
  • Route pure or mixed IQ MoE to batched MMQ from M=128 only when assignment
    reuse, expert count, CUDA grid capacity, and both projection K alignments are
    suitable.
  • Add an AOT capability query and enable IQ MMQ in SRT only when the installed
    sglang-kernel exposes that capability.
  • Reject unsupported Q8_1 in the high-level selector and direct CUDA GGUF
    compute entry points.
  • Harden shared MoE contracts with scalar catch-all registration, 64-bit expert
    strides, padded-tail bounds, complete q8 scale-tile loading, low-level
    alignment checks, and deterministic zero output for invalid or remote
    experts.
  • Add registered CPU routing tests, CUDA dense and pure/mixed MoE regression
    coverage, and a checked-in dense benchmark using a revision-pinned GGUF
    fixture.

The current SGLang dependency remains pinned to
sglang-kernel==0.4.6.post1. That released wheel does not expose the new
capability, so the patched SRT safely retains its existing dequant/vector
fallbacks. After this kernel source lands, a new sglang-kernel wheel must be
published and the SGLang pin updated in a follow-up before normal installations
enable the new MMQ path. This PR intentionally does not bump the pin to an
unreleased wheel.

The kernel implementation is adapted from
aimbit-ni/vllm@037e1d547c15313aa7e2bb5fc83390ab0c857314 and the algorithms
in ggml-org/llama.cpp#8495. THIRDPARTYNOTICES.txt includes the ggml MIT
notice.

Accuracy Tests

Local validation used the installed CUDA 13/aarch64 AOT wheel on one NVIDIA
GB10 (sm_121), rather than JIT-compiled kernels.

The main-sync validation pair was upstream base
0c84eaba7f8d6674a360839426cacde4e1acd851 and clean local merge head
4c3e4c2350f4d6c4ff2a785e9b294938766a3cc6, which is now the remote Draft PR
head.

  • validated source-equivalent wheel SHA-256:
    626d6ee79dd37dae2242f2f18a97b0536c95dfa4c053cb56d7573de23324c6e0
  • registered CPU routing and compatibility suite: 10 passed
  • patch-focused GGUF CUDA matrix: 57 passed
  • complete python/sglang/kernels/aot/tests/test_gguf.py:
    153 passed
  • installed-wheel production smoke: passed

Using the same CUDA 13/aarch64 build configuration, the compressed wheel grew
from 29.473 MiB to 30.999 MiB: +1.527 MiB (+5.18%). Each bundled
architecture-specific common_ops shared object grew by 917,504 bytes.

The wheel was built before the final main-only merge. The AOT tree, SRT GGUF
selector, and registered routing test have identical Git object IDs at the
wheel-build head and the validation head above; the intervening upstream commits
do not change this source or test surface.

The production smoke observed the new capability as enabled and covered:

  • wide dense IQ MMQ at M=16 and dequantization fallback at M=17
  • exact fixed-shape CUDA Graph capture/replay
  • MoE vector dispatch at M=127 and batched dispatch at M=128
  • pure and mixed quantization at M=128, M=256, and M=2050
  • Q8_1 rejection, K-alignment guards, and invalid-expert zero output

Maximum reserved GPU memory during the production smoke was 528 MiB.

A compatibility smoke combined the patched SRT source with the currently pinned
sglang-kernel==0.4.6.post1 baseline wheel (SHA-256
c7aac3ab20ff3c704e2bbc1a2fa78e5142a65150f12bca74f8e4cbd4f651d958).
It observed capability=False, selected dequantize for an IQ dense input, and
returned shape [9, 18944]. This verifies that the combined source remains
correct before the new kernel wheel is released.

A final-head bounded end-to-end serving sanity check also passed with the
artifacts in /tmp/sglang-36122-e2e-final.XwQYAu and the revision-pinned
legraphista/Qwen2.5-0.5B-Instruct-IMat-GGUF
Qwen2.5-0.5B-Instruct.IQ3_S.gguf file (SHA-256
b7685f565fc6e49a8f9bc0fe09b9bb22ca0d5e4b10635cd4b181597f113794b6).
The installed patch wheel (SHA-256 626d6ee7...) reported capability=True;
the same synced SRT with the released baseline wheel (SHA-256 c7aac3ab...)
reported capability=False. Both configurations loaded the model, returned
HTTP 200 from the health and generation requests, and recorded zero cached
tokens for generation. Their client JSON files were byte-for-byte identical
with SHA-256
bd3c11d390f7a373e534f5361807119b2c040bc325cfe1f6fd2c8165d4426933:
11 prompt tokens, 16 completion tokens, a length finish, and text:

 I like to work outside because it is a great way to get exercise. I

Both runs used a single GB10 (sm_121, CUDA 13, aarch64), read-only source,
wheel, and model mounts, no network or host port, a 10 GiB container memory
limit, --mem-fraction-static 0.03, --max-total-tokens 1024, one running
request, Triton attention, PyTorch sampling, and CUDA graphs disabled. The patch
run loaded 0.88 GB of weights and allocated approximately 0.02 GB of KV cache.
This validates a small dense IQ3_S model and the released-wheel fallback; it
does not replace testing the original approximately 100 GB routed-expert model.
The critical server-log scan found no traceback, CUDA error, NaN, or
illegal-memory failure; the remaining messages were the expected nvidia-smi
fallback and shared local-tokenizer warnings.

The official Lint workflow
passed for the prior remote head 0137a1e..., including
pre-commit run --all-files, the repository lint checks, and the dedicated
sgl-kernel clang-format check. A new Lint run for the synced head is in progress;
normal platform CI remains Draft-gated.

The complete local AOT collection is not claimed as passing. A bounded run
reached 971 passed and 2 skipped before an unrelated large AWQ parameter
exceeded the allocator cap. Excluding AWQ reached 970 passed and 2,310 skipped
before FlashMLA sparse prefill rejected SM121.

Normal/full PR CI is still pending while this PR is a Draft. The current red PR
Test entries are draft-gate failures and their wheel, H100, and B200 jobs were
skipped; they are not kernel test failures. Before merge, the PR must be marked
ready and run the normal full workflow so it builds and installs the PR-local
wheel on the official x86_64/aarch64 and H100/B200 environments.

Speed Tests and Profiling

These are directional route-selection microbenchmarks from a shared GB10 that
retained a live Qwen service. They are not uncontended end-to-end throughput or
cross-architecture claims. They predate the main sync and were not rerun; the
sync did not change the IQ kernel patch represented by these measurements.

The checked-in benchmark compares direct MMQ with dequantization plus BF16
matmul using revision-pinned 18944x1024 GGUF fixtures. The table below uses
two independent benchmark invocations. Each value is the mean speedup, with the
two-run range in brackets; speedup is (dequantize + BF16 matmul) / MMQ.

quant type M=16 speedup, mean [range]
IQ1_S 1.75x [1.75x, 1.76x]
IQ2_XXS 1.31x [1.29x, 1.33x]
IQ2_XS 1.26x [1.26x, 1.26x]
IQ2_S 1.60x [1.59x, 1.61x]
IQ3_XXS 1.71x [1.67x, 1.74x]
IQ3_S 1.66x [1.65x, 1.67x]
IQ4_NL 1.48x [1.47x, 1.48x]
IQ4_XS 1.26x [1.24x, 1.28x]

The full M=9/16/17/24/32 matrix and reproduction notes are in the
benchmark comment.
The M=16 ranges retain at least 1.24x measured headroom across all eight
formats. Although the isolated M=17 sample was also faster, M=18..23 were not
measured and M=24 regressed on multiple formats. The shared GB10 measurements
do not justify interpolating a larger universal CUDA gate, so this PR uses the
conservative M<=16 boundary. M=17 and above retain dequantization.

Directional MoE measurements were:

path vector baseline batched MMQ speedup
IQ3_S/Q5_0 M=128 2.249 ms 1.490 ms 1.51x
Q5_0/IQ3_S M=128 2.524 ms 1.571 ms 1.61x
IQ3_S/Q5_0 M=256 4.497 ms 1.758 ms 2.56x
Q5_0/IQ3_S M=256 5.145 ms 2.283 ms 2.25x
IQ3_S/Q5_0 M=2050 34.801 ms 9.056 ms 3.84x

TTFT, TPOT, serving throughput, and the original approximately 100 GB report
model remain unmeasured. The representative serving check above is an accuracy
and compatibility sanity check, not a throughput result.

The Issue author has offered to run the original model and benchmark harness as
a pinned three-arm 4k/16k comparison: the upstream snapshot with the released
wheel, the refreshed PR source with the released wheel, and the same PR source
with a wheel built from that checkout. That external result is still pending,
so the PR remains Draft; see the
reviewer response
and the
pinned U/H protocol.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #32939745573
Latest PR Test (Extra): ❌ Run #32939745211
Latest PR Test (AMD ROCm 7.2): ❌ Run #32939745576

Add dense and batched MoE MMQ kernels for eight supported I-quant GGUF formats, with conservative routing and compatibility fallback for older sglang-kernel wheels.

Fixes sgl-project#35019

Signed-off-by: Hao Liu <liuhao24@sjtu.edu.cn>
@liuhao-labs

liuhao-labs commented Aug 24, 2026

Copy link
Copy Markdown
Author

GB10 dense IQ MMQ crossover

GGUF IQ MMQ crossover on GB10

This plot uses the revision-pinned 18944x1024 GGUF fixtures on CUDA 13 / GB10
(sm_121). Each cell is the mean speedup from two independent invocations of
the checked-in benchmark, with the two-run range shown below it. Speedup is
(dequantize + BF16 matmul) / MMQ.

At the sampled MMQ boundary M=16, the two-run ranges retain at least 1.24x
measured headroom across all eight formats. The isolated M=17 sample was also
faster, but M=18..23 were not measured and M=24 regressed on multiple formats.
To avoid interpolating a universal CUDA threshold from one shared GB10, the PR
uses the conservative M<=16 gate; M=17 and above remain on dequantization.

Run the underlying benchmark twice from python/sglang/kernels/aot in an
environment with the current PR wheel installed. The --local-files-only flag
requires the pinned fixtures to already be present in the Hugging Face cache;
omit it on the first run if the environment has network access.

for run in 1 2; do
  python benchmark/bench_gguf_mmq.py \
    --batch-sizes 9,16,17,24,32 \
    --shape 18944x1024 \
    --local-files-only | tee "/tmp/gguf-mmq-run-${run}.md"
done

The heatmap was rendered from those two benchmark tables; the plotting-only
helper is not part of this PR. These are directional local microbenchmarks on a
shared GB10 with a live Qwen service, not uncontended end-to-end throughput or
cross-architecture results.

@liuhao-labs

Copy link
Copy Markdown
Author

Additional end-to-end compatibility check

I ran the revision-pinned
legraphista/Qwen2.5-0.5B-Instruct-IMat-GGUF
Qwen2.5-0.5B-Instruct.IQ3_S.gguf file (SHA-256
b7685f565fc6e49a8f9bc0fe09b9bb22ca0d5e4b10635cd4b181597f113794b6)
through SGLang's real /generate endpoint on a GB10 (sm_121, CUDA 13,
aarch64).

configuration capability prompt/completion tokens result
patched SRT + PR-local wheel (dae0fcc2...) true 11 / 16 pass
patched SRT + released sglang-kernel==0.4.6.post1 (c7aac3ab...) false 11 / 16 pass

Both greedy requests produced byte-for-byte identical output:

 I like to work outside because it is a great way to get exercise. I

The runs used read-only source/wheel/model mounts, no network or host port, a
10 GiB container limit, --mem-fraction-static 0.03,
--max-total-tokens 1024, one running request, Triton attention, PyTorch
sampling, and disabled CUDA graphs. The patch run loaded 0.88 GB of weights and
allocated approximately 0.02 GB of KV cache.

This is a small dense IQ3_S serving and released-wheel fallback check. It does
not claim end-to-end coverage of the original approximately 100 GB
routed-expert model; pure and mixed MoE paths are covered by the focused CUDA
and installed-wheel production tests described in the PR body.

@liuhao-labs

Copy link
Copy Markdown
Author

Refreshed-head validation

The Draft branch is now synchronized and pinned to:

  • U = 0c84eaba7f8d6674a360839426cacde4e1acd851 — upstream main snapshot used as the validation baseline
  • H = 4c3e4c2350f4d6c4ff2a785e9b294938766a3cc6 — current Draft PR head; this merge commit contains U

The feature diff from U to H remains 15 files (+4127/-81). The CUDA 13/aarch64 wheel SHA-256 is 626d6ee79dd37dae2242f2f18a97b0536c95dfa4c053cb56d7573de23324c6e0 (30.999 MiB). The wheel-build head and H have identical AOT tree, SRT GGUF selector, and registered routing-test Git objects; the final merge only incorporated upstream changes outside those validated surfaces.

Validation on one GB10 (sm_121, CUDA 13, aarch64):

  • registered CPU routing/compatibility: 10 passed
  • focused GGUF CUDA matrix: 57 passed, 5 warnings in 49.07s
  • complete test_gguf.py: 153 passed, 5 warnings in 60.32s
  • installed-wheel production smoke: passed with a 0.06 allocator cap (528 MiB maximum reserved)
    • dense M=16 -> sgl_kernel::ggml_mul_mat_a8
    • dense M=17 -> sgl_kernel::ggml_dequantize
    • MoE M=127 -> sgl_kernel::ggml_moe_a8_vec
    • MoE M=128 and M=2050 -> sgl_kernel::ggml_moe_a8
    • fixed-shape CUDA Graph capture/replay matched exactly

I also reran the small revision-pinned IQ3_S serving check on H with two kernel-wheel pairings:

pairing capability health/generate cached prompt tokens result
H source + refreshed PR wheel (626d6ee7...) true 200 / 200 0 pass
H source + released sglang-kernel==0.4.6.post1 (c7aac3ab...) false 200 / 200 0 pass

The two client JSON files were byte-for-byte identical (SHA-256 bd3c11d390f7a373e534f5361807119b2c040bc325cfe1f6fd2c8165d4426933): 11 prompt tokens, 16 completion tokens, length finish, and identical greedy text. The critical server-log scan found no traceback, CUDA error, NaN, OOM, or illegal-memory failure.

This remains a Draft. The original approximately 100 GB routed-expert 4k/16k A/B, official H100/B200 AOT runs, and official x86_64/aarch64 wheel CI are still pending.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Perf] No MMQ kernel for I-quant GGUF: 4-6x slower prefill than llama.cpp on the same weights (measured)

1 participant