Skip to content

ryanmurf/qwen-kernel

Repository files navigation

qwen-kernel

A from-scratch Vulkan inference engine + serving stack for Qwen3.6-35B-A3B on RDNA3 (Radeon RX 7900 XT). One model, fully specialized: hand-written compute kernels for every weight format in the GGUF, the whole hybrid gated-DeltaNet + MoE architecture fused into pre-recorded command buffers, and a safe-Rust server speaking the Anthropic Messages API — so Claude Code runs against it directly. Inspired by KernelBench Mega (CUDA-only); this is the RDNA3/Vulkan equivalent, taken all the way to a serving engine.

Hardware scope: tested only on RDNA3 (Radeon RX 7900 XT / XTX) with Mesa RADV on Linux. The kernels assume RDNA3 subgroup behaviour. It will build on other vendors — Vulkan is Vulkan — and then not work correctly. There is also an experimental Apple/Metal backend (src/main_metal.mm, shaders/metal/); it is not the primary target.

Benchmarks

End-to-end through the Claude Code CLI (same GGUF, same GPU, one backend at a time — full methodology and raw data in bench/, 2026-07-08):

scenario (median wall) qk-server llama.cpp Vulkan (master) qk
single-shot generate ~350 tok 5.3 s 32.5 s 6.1×
Bash tool call (2-turn agentic loop) 1.5 s 6.3 s 4.3×
multi-turn resumed session, per turn 2.3 s 14.3 s 6.2×
warm decode through the CLI 65 tok/s 12 tok/s 5.4×

Engine-level, token-exact greedy parity with llama.cpp verified across the full stack. Decode re-measured 2026-07-18 against llama.cpp 571d0d5 (same day), same model file, f16 KV both sides — see bench/:

metric qk llama.cpp Vulkan 571d0d5
single-stream decode, 7900 XTX 190.7 tok/s (5.24 ms/tok) 132.3 tok/s (1.44×)
single-stream decode, 7900 XT 147.1 tok/s (6.80 ms/tok) 109.7 tok/s (1.34×)
prefill pp512, 7900 XTX 2904.2 tok/s 2900.3 tok/s (1.00× — parity)
prefill pp1024, 7900 XTX 3098.1 tok/s 2852.0 tok/s (1.09×)
prefill pp2048, 7900 XTX 2907.9 tok/s 2812.6 tok/s (1.03×)
aggregate decode, 16 streams 384 tok/s
warm start from prefix cache 0.3 ms restore (vs 341 ms/64-tok prefill)

Prefill reached parity on 2026-07-19 after a cooperative-matrix (KHR_coopmat) rewrite of the DeltaNet batch step, attention core, dense projections and router-logits GEMM. Note pp512 is parity, not a win — 1.001× sits inside llama.cpp's own ±1.4% run-to-run spread. pp1024 and pp2048 are genuine wins. Both sides are measured hot (repetitions back-to-back with gpu_busy verified 0 beforehand), which is llama-bench's protocol; qk's earlier published prefill numbers used a polling harness that forced a ≥100 ms idle between repetitions and therefore paid an ~11% down-clocked-DPM penalty llama.cpp never paid.

† Not re-measured on 2026-07-18 and not backed by raw data in bench/; treat as provisional. The end-to-end CLI table above uses a llama.cpp build from 2026-04-05 and overstates the current gap — see bench/README.md.

Correctness bar throughout: greedy output is token-for-token identical to llama.cpp on identical input ids, batched paths are validated bit-identical (or argmax-stable at ~1e-7 rel) against serial references, and the server's tokenizer reproduces llama.cpp byte-for-byte.

How it works

  • src/main.cpp + shaders/ — the engine (libqk.so / qk CLI). GEMV/GEMM kernels for Q8_0, Q6_K, IQ4_XS, IQ3_XXS, F16 at 90–97% of VRAM bandwidth on the big formats; the fused MoE step (256 experts, top-8 + shared) and gated-DeltaNet recurrence (state resident on GPU); full attention with GQA + partial NeoX rope; GPU-resident argmax sampling. A whole decode step is pre-recorded command buffers — one queue submit per chunk, host reads ids at the end. N slots batch on the dispatch z-axis, so concurrent requests of different lengths share every weight read.
  • server/qk-server, a safe-Rust (axum) HTTP layer over the one dlopen'd engine thread: llama.cpp-compatible endpoints plus native Anthropic POST /v1/messages (SSE streaming, hermes <tool_call> parsing, Qwen3 chat template), a GGUF-native BPE tokenizer, per-slot admission, context-fit trimming with fast-fail. See docs/server-spec.md.
  • Prefix / cross-turn KV reuse — a prompt's KV + recurrent state is snapshotted at the conversation-history boundary and restored on the next turn, so a growing agentic session prefills only each turn's delta (O(N) per session instead of O(N²)). Production hit-rate ~88%; agentic turns ~5× faster. QK_PCACHE_LOG=1 prints per-request reuse stats.
  • Batched prefill — prompt chunks run as one command buffer with a register-blocked Q8_0 GEMM (weights read once per chunk instead of once per token) and the DeltaNet recurrence collapsed into one workgroup-per-head dispatch. Auto-selected; QK_NO_BATCH=1 forces the serial reference path. Chunk width saturates at 128 (measured; 256 is token-exact but +1%).
  • Pipeline split (QK_LAYERS=a:b) — an engine instance can own just transformer layers [a,b): the first stage also owns the embedding, the last owns final-norm + head + argmax, and the ~8 KB/token residual row crosses stage boundaries through the qk_stage_run ABI (in-process or TCP — qk pipe / qk pipe-worker). Greedy output is token-exact vs the unsplit engine at every tested boundary; two stages of a 20-layer half hold ~8.5 GB each vs ~15.4 GB whole. Foundation for serving models larger than one GPU (e.g. tron + midnight once the Metal port lands).
  • deploy/ — k8s deployment sharing one GPU with a llama.cpp fallback backend (switch.sh qk|gemma), image build script, and an orphaned-GPU-process reaper. See deploy/README.md.

Model facts (qwen35moe): n_embd 2048, 40 blocks (30 gated-DeltaNet + 10 full-attention), 256 experts top-8 + 1 shared, vocab 248,320. The UD-Q3_K_M GGUF actually contains Q8_0 / Q6_K / IQ4_XS / IQ3_XXS tensors — no Q3_K. The engine is deliberately hard-wired to this architecture; that specialization is where the speed comes from.

Build & run

Prerequisites

  • CMake >= 3.16
  • Vulkan headers + loader (libvulkan-dev / vulkan-headers)
  • glslc (from shaderc) — required to compile the shaders
  • a C++17 compiler
  • Rust toolchain supporting edition 2024 (recent stable) — for the server only
  • Mesa RADV. Tested on Mesa 25.x with a 7900 XT / XTX.
  • spirv-val (from spirv-tools) if you want to validate shaders
cmake -B build
cmake --build build -j

If cmake on your system is a broken pip shim (it happens), use /usr/bin/cmake explicitly.

# serve (Anthropic + llama.cpp-compatible HTTP)
cd server && cargo build --release
QK_SHADER_DIR=../build/shaders ./target/release/server \
    --model /path/Qwen3.6-35B-A3B-UD-Q3_K_M.gguf \
    --engine-lib ../build/libqk.so --port 8080 --slots 2 --ctx 32768 --chunk 8

Engine CLI (./build/qk …) — benchmarks and correctness harnesses:

qk                       # synthetic kernel suite (f16, q8_0, q6_k, iq4_xs, iq3_xxs)
qk gguf <tensor>         # GEMV on real weights          qk moe|block|ablock [layer]
qk token <ids> <n> [tmax] [batch]   # end-to-end greedy generation
qk warm <ids> <n>        # prefix-cache TTFT demo        qk serve-test <ids> <n> [slots]
qk prefillcmp|prefillbench|prefilldecode   # batched-prefill exactness / timing / handoff
qk verify <ids> <n> [K,..]   # spec-decode verify rounds (oracle draft): exactness + c(K)
qk pipe <ids> <n> [split] [tmax] [host:port]   # pipeline-split parity/timing (2 stages)
qk pipe-worker <port> [a:b] [tmax] [slots]     # serve one stage over TCP
qk list [filter]         # tensors in the GGUF

Env knobs:

var meaning
QK_GGUF, QK_DEVICE, QK_SHADER_DIR model path, Vulkan device index, SPIR-V dir
(QK_GGUF/--model also accept the first shard of an llama.cpp-style split model, …-00001-of-NNNNN.gguf)
QK_FORK=1 prefix cache: same-prefix requests restore instead of re-prefilling
QK_LAYERS=a:b pipeline split: this engine owns layers [a,b) only, driven via qk_stage_run
QK_SPEC=1 prompt-lookup speculative decoding (exact output; ~1.5× on echo-heavy generation)
QK_SPEC_K, QK_SPEC_L, QK_SPEC_LOG=1 verify width (8), trigger n-gram length (6), per-request [spec] stats
QK_PCACHE, QK_PCACHE_LOG=1 prefix-cache LRU depth (default 3; ~1.31 GiB host RAM per entry at ctx 32768), per-request stats
QK_STATS_FILE also append [spec]/[pcache] lines to this file (survives pod restarts)
QK_NO_BATCH=1 force serial prefill (correctness reference)
QK_MAXB batch-prefill chunk width (default 128; buffers scale with it)
QK_SUBMIT_LAYERS, QK_ATTN_BUDGET submit granularity / attention tile budget (amdgpu ~10 s ring-timeout guards)
QK_MAX_TOOL_CHARS server: cap on a single tool_result before context-fit trimming

Docs

Notes

  • Kernel-bench numbers want a quiet GPU: scale the serving pod down first (kubectl scale deploy qk-server -n gemma --replicas=0, and back to 1 after). The engine needs ~15.4 GB resident for end-to-end generation.
  • Development history — the M1→M6 bring-up (GEMV → quantized kernels → fused MoE → DeltaNet blocks → attention → end-to-end), the optimization arc 33.3 → 5.82 ms/token, and measured negative results (IQ3 repacking, RoPE precompute) — lives in this file's git history and the commit log.
  • Reference shaders: llama.cpp's Vulkan backend (ggml/src/ggml-vulkan/vulkan-shaders/).

Why is it faster (and where it isn't)

Every compute kernel here is hand-written. From llama.cpp/ggml this project takes only the quantization lookup tables in shaders/iq_tables.glsl, the work shape of three Metal GEMVs, and the GGUF container format — see NOTICE. No llama.cpp shader or binary is compiled, linked or executed.

The decode advantage comes from three things, roughly in order of size:

  1. Pre-recorded command buffers. llama.cpp walks a generic graph and dispatches per node; qk records a whole decode step once and submits it as a single command buffer, reading token ids back only at the end. At batch 1 that removes hundreds of tiny dispatches per token. (llama.cpp has an equivalent for CUDA via GGML_CUDA_GRAPHS, but its Vulkan backend has no command-buffer-reuse path.)
  2. Shape specialization. Kernels are compiled for the exact K, n_ff and expert count, so trip counts and tails are compile-time constants and there is no runtime shape dispatch. A general runtime cannot do this.
  3. Architecture-specific fusion. The gated-DeltaNet recurrent state stays resident on the GPU, route+select is fused, expert dispatches are grouped, and norm/residual epilogues are folded into their producers.

Prefill is a different regime: matrix cores only help when there is data reuse. Batch-1 decode is a matrix-vector product — each weight is read once and discarded, so it is bandwidth-bound and RDNA3's KHR_coopmat units cannot help (our Q4_0 GEMV already runs at ~87% of achievable bandwidth). Prefill is a matrix-matrix product where 512 tokens share every weight read; it is compute-bound, and coopmat is exactly the right tool. This engine originally used none there and lost prefill ~2.4×; adding native cooperative-matrix kernels closed it to parity.

The instructive part was where the time actually went. Three optimization passes attacked the routed-MoE gate/up kernel on the assumption it dominated. Dumping llama.cpp's own per-op times (GGML_VK_PERF_LOGGER) and lining them up against qk's stage counters showed qk was already faster there (59.0 vs 63.4 ms). The real deficits were the DeltaNet batch step (32.4 vs 6.3 ms), the attention core (17.1 vs 1.6 ms), and dense projections held on a scalar path by a precision constraint. Profiling the opponent, not just ourselves, is what found it.

Porting to other hardware: the ideas in (1)-(3) are architecture-independent, but every tile size, threads-per-row and split width here was chosen by sweep on 96 CUs with 96 MB of Infinity Cache. Expect to re-tune all of them — in related work on this repo the best threads-per-row varied from 16 to 256 across shapes on the same card.

Known issues

  • tests/ids4.txt diverges from ref4.txt at token 32 (735 vs 13914). This is pre-existing and reproduces on an untouched tree; it is not yet diagnosed. tests/ids2.txt and ids4.txt also produce produced=0 under serve-bench, and block/ablock carry pre-existing Vulkan CPU-reference threshold failures. Given that token-exact parity is this project's headline claim, these are being treated as real and are under investigation.
  • The 80B path requires a repacked -qk GGUF produced by a tool that has not been published, so it is not reproducible outside this machine.

License

MIT — see LICENSE.

This project derives from llama.cpp / ggml (also MIT): the quantization lookup tables in shaders/iq_tables.glsl, the Metal kernel work shape for three GEMVs, and the GGUF quantization block layouts. See NOTICE for details. llama.cpp is also the reference implementation for all token-exact parity testing here.

Model weights are not distributed with this software and carry their own licenses from their respective publishers.

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages