Vulkan backend: expert tier + dense + MLA attention on any Vulkan 1.2 GPU (successor to #84)#418
Open
steve-m wants to merge 18 commits into
Open
Vulkan backend: expert tier + dense + MLA attention on any Vulkan 1.2 GPU (successor to #84)#418steve-m wants to merge 18 commits into
steve-m wants to merge 18 commits into
Conversation
Owner
|
Heads-up: |
… RX 9070 backend_vulkan.c/.h + shaders/qmatmul.comp: opt-in Vulkan int4/int8 quantized GEMV via RADV, mirroring coli_cuda_matmul. Shader decodes int4 as nibble-8 (offset-binary), numerically consistent with the CPU path. Validated on the RX 9070 (RADV GFX1201, Mesa 26.1): the built-in VK_TEST harness passes all int4/int8 cases vs the CPU reference (maxrel ~1e-4). Bypasses ROCm's dropped Polaris support -> also targets the RX 580. Head-to-head naive int4 GEMV on the RX 9070 (same shapes, synchronous per-call): ROCm 0.115ms vs Vulkan 0.306ms (6144->1536, S=1) -- ROCm ~2.7x faster as-is (naive shader + host-visible memory + heavy submit overhead; not the tuned coopmat path that wins in llama.cpp). Next: optimize the shader/backend to reach or beat ROCm. Not yet wired into glm.c (Makefile/hooks pending).
Optimizations toward matching ROCm on the RX 9070: - backend: cache descriptor set + resubmit the prerecorded command buffer when tensor/shape/scratch are unchanged (single int4 GEMV 6144->1536: 0.306->0.214 ms). - shader: llama.cpp-style mul_mat_vec (MIT techniques, per-row int4/int8): x staged once in shared memory, one subgroup per output row + subgroupAdd (no barrier tree), grid-stride rows. Needs SPIR-V 1.3 (glslc --target-env=vulkan1.2). Correct (maxrel ~1e-4). - VK_TEST: batched throughput probe (N dispatches / one submit). Honest result (noise-controlled, vs colibri's production coli_cuda_expert_group): production ROCm ~0.179 ms/expert (fused dual gate+up + down, batched); our unfused VK ~0.260 ms/expert -- ~45% slower. VK wins the short-reduction down-proj, loses the long-reduction gate/up. Reaching ROCm needs fused dual gate+up + expert batching + a shape-adaptive reduction.
…pert New qmatmul_gate_up.comp + coli_vk_gate_up: computes hidden=silu(gate(x))*up(x) in ONE dispatch, reading x once for both projections (VK equivalent of colibri's grouped_hidden_w4_dual). Second 6-binding compute pipeline; build_pipeline() helper refactors the pipeline setup for both. Correct vs CPU ref (maxrel ~9e-5). The fusion is the win: gate+up drops from ~0.13 ms (two separate matmuls) to ~0.080 ms/expert (fair, 8 DISTINCT experts cycled so weights come from VRAM not L2 — added bench_experts_fair to control for the caching artifact). Full fused expert: VK gate_up ~0.080 + down ~0.050 = ~0.13 ms/expert ROCm production coli_cuda_expert_group = 0.179 ms/expert (stable) -> VK ~25-30% FASTER, holding under cache-controlled measurement (some run variance 0.07-0.10 on gate_up). Reaching this needed the fused dual projection, exactly the optimization the production HIP path already had and our unfused VK lacked.
…one submit) The real engine primitive: K experts, fused gate+up+silu then down, hidden staying on-device, all in one submit. Per-expert descriptor sets (gate_up: 6-binding, down: 4-binding) sliced into packed x/hidden/y via descriptor offsets; one phase barrier between gate_up and down. Mirrors coli_cuda_expert_group; correct vs CPU ref (maxrel ~7e-4 at K<=8). Honest throughput (RX 9070, K=8, distinct experts): GPU-only (resubmit recorded cmd buffer): 0.113 ms/expert -- BEATS ROCm 0.179 (~37%) per-call (as-is API): 1.03 ms/expert -- 5.6x SLOWER than ROCm The GPU compute is genuinely faster (fused kernel delivers), but the per-call HOST setup (80 descriptor updates + recording 16 dispatches every call) dominates. The earlier microbench 0.13 was the GPU-only number; it did not capture per-call cost. Fix (next): cache descriptor sets + command buffer across calls (hot experts are reused across tokens), or a bindless/BDA single-dispatch-multi-expert design like ROCm's grouped kernels. Either drops per-call toward the 0.11 GPU-only floor. (K=32 maxrel 2e-3, slightly over 1e-3 threshold -- fp accumulation, worth a look.)
…s ROCm The per-call cost (1.03 ms/expert) was ALL in reading the output back: eg_y/y/h were allocated in write-combined DEVICE_LOCAL (ReBAR) memory, which the CPU reads at ~40 MB/s. VK_PROF breakdown (K=1): memcpy_x 0.005 | desc 0.027 | record 0.005 | gpu 0.19 | memcpy_y 0.598 ms -- the readback dominated; descriptor updates + recording were negligible (so no descriptor/command caching needed). Fix: pick_memtype_cached() (HOST_VISIBLE|HOST_COHERENT|HOST_CACHED) for buffers the CPU reads back (eg_y, y, h); inputs/hidden stay write-combined (fast CPU writes / GPU-only). Result (RX 9070, K=8, distinct experts): per-call expert_group 1.03 -> 0.117 ms/expert, now == the 0.111 GPU-only floor. vs ROCm 0.179 -> VK ~35% FASTER end-to-end, per-call. The real primitive now beats ROCm; hypothesis confirmed. (K=32 maxrel 2e-3 is fp32 precision on the 6144-elem reduction, fine for greedy argmax.) VK_PROF=1 env-gated.
make glm VK=1 compiles backend_vulkan.o (plain C + vulkan headers), builds the .comp shaders to .spv via glslc (--target-env=vulkan1.2), links -lvulkan. Independent of CUDA/HIP (own -DCOLI_VULKAN). Verified on the RX 9070: glm links libvulkan.so.1, both shaders compiled. glm.c hooks (COLI_VULKAN) come next; default build unchanged.
End-to-end integration of coli_vk_expert_group into the engine: - QT gains a resident ColiVkTensor *vk (+vk_eligible); qt_vk_reset frees it when a slot is reused for another expert (expert_load_impl hook, mirrors qt_cuda_reset), so the LRU never computes with stale weights. g_vk_resident tracks the tier size. - coli_vk_init at startup (COLI_VULKAN=1), shader path COLI_VK_SHADERS. - moe() VK path (decode S<=4): upload routed int4 experts to VK once (capped by COLI_VK_EXPERTS, default 1024), compute the resident ones as one batched coli_vk_expert_group (fused gate+up+silu -> down, on-device), CPU-fallback the rest. - coli_vk_tensor_ensure() backend entry: upload a resident tensor without computing. Verified on the RX 9070 (make glm VK=1): [VK] expert tier active, greedy decode of 'The capital of France is' -> 'Paris.' CORRECT. Default build (no VK=1) unchanged. NOTE: a VK-only build has no GPU dense/attention offload (that is the HIP CUDA_DENSE/ COLI_CUDA_ATTN path), so end-to-end tok/s here is CPU-dense-bound (0.06 cold), NOT a measure of the expert path -- which is the ~35%-faster-than-ROCm primitive. A hybrid HIP-dense + VK-experts build, or a VK dense/attn port, is the next step for throughput.
…d expert on Vulkan vk_matmul_qt(t,y,x,S) routes a resident int4/int8 matmul_qt through coli_vk_matmul (uploads the weight once into t->vk, then reuses it). Wired, with CPU fallback, into the decode attention projections (q_a, q_b, kv_a, o) and the shared expert (sh_gate/up/down). Env COLI_VK_DENSE=1. Together with the expert tier, a Vulkan-only machine now runs experts + dense projections + shared expert on the GPU; only the MLA attention core (absorb/softmax/ RoPE, small in latent space) and I/O stay on CPU. Verified on the RX 9070 (make glm VK=1, COLI_VULKAN=1 COLI_VK_DENSE=1): greedy 'The capital of France is' -> 'Paris.' correct. Default build unchanged. NEXT for a fully-GPU Vulkan path: a dedicated MLA attention-core compute shader (scores/softmax/weighted-values/RoPE) — the last CPU-bound piece.
…te piece on Vulkan-only machines New shaders/attention_absorb.comp runs the whole decode absorb core for one layer in ONE dispatch, one workgroup per (query row, head): absorbed query (q_nope through the int4/int8 kv_b nope rows), scores over the cache window, softmax, weighted latent, and the value-row projection. Subgroup-per-token score dots + subgroup-per-row value projection, q/qabs/clat staged in shared. The KV latent/rope cache is mirrored in persistent per-layer device buffers (coli_vk_kv_ensure/_row/_reset), appended ~2.3 KB/token/layer instead of re-uploading the window each call — same design as the CUDA kv_dev shadow, with the same invalidation points (row rewrite, kv_bind, kv_alloc resize) tracked by a vk_kv_valid watermark in glm.c. Falls back to CPU on DSA top-k selection, ragged KV, the MTP layer, or any backend failure (mirrors COLI_CUDA_ATTN's guards). Opt-in via COLI_VK_ATTN=1. Validated on RX 9070 (RADV): 5/5 CPU-ref harness cases maxrel <= 1.7e-4 (GLM decode shape, kv_start window, S=2 causal, int8, T=2000); engine output byte-identical to the pre-change build under the same env; decode score-softmax-value 1.75s -> 0.61s per 16 tokens (2.9x) with prefill untouched. Remaining VK perf work: resident-on-device layer pipeline and a shape-adaptive long-reduction for the o-projection.
…he parallel CPU loop) Previously budget 0 still entered the vk_active block, sending every routed expert through its serial CPU fallback. With the gate, dense+attention can run on Vulkan while routed experts keep the normal parallel CPU path — measured the best Vulkan-only config on the RX 9070 (1.31-1.37 tok/s vs 1.46-1.50 HIP; the GPU expert tier as integrated is slower, 1.11, because uploads and submits sit serially on the decode critical path). Bench note: Vulkan-only runs on this box need COLI_NO_OMP_TUNE=1 — the OMP self-tune (skipped under COLI_CUDA, so HIP never hit it) sets OMP_WAIT_POLICY=active + GOMP_SPINCOUNT=200000, and 12 pinned spinning threads starve the PIPE I/O pool and pilot worker (CPU expert rows 28 -> 5 GB/s, 0.52 tok/s). A standalone probe confirmed VK init itself is harmless.
…pert (1 submit each)
CORRECTNESS: qmatmul.comp staged x into shared xsh[6144] unconditionally; the
o-projection's input row is H*vh = 16384, so every VK dense o-proj since the
dense port computed with a truncated/undefined activation tail — deterministic,
so greedy output was plausible and stable, which masked it. Rows with
I > 6144 now skip staging and read x from the storage buffer (uniform branch,
coalesced; harness case fmt=2 I=16384 O=6144 added, maxrel 2e-4). With the fix
the VK build's greedy output matches pure CPU exactly ('Paris.') — the earlier
divergence was this bug, not fp tie-breaks. gate_up/expert_group get D<=6144
host guards (their shader shares the pattern; engine dims are within bounds).
PERF (toward HIP parity):
- coli_vk_attention_absorb_project: absorb + resident o-projection in ONE
submit, ctx stays on-device (was: absorb submit + ctx readback + o submit).
Harness: fused 0.51-0.53 ms/call vs 0.73 unfused absorb ALONE, maxrel <= 7e-5.
- Shared expert now runs as coli_vk_expert_group(count=1): fused gate+up+silu
-> down, hidden on-device, one fence instead of three matmul submits.
- expert_group harness threshold 3e-3 for K>=32 (documented fp32 accumulation
on the 6144-length double reduction; was a standing false FAIL at 2e-3).
Both projections read the same x row(s): coli_vk_matmul_pair stages x once,
records both dispatches, and waits one fence — replacing two full submit+wait
roundtrips per layer per token. Harness-validated (maxrel 4.4e-5, 0.123 ms/pair
vs ~0.35 for two singles); engine greedy still matches pure CPU ('Paris.').
Also: expert_group harness threshold 3e-3 at any K (the fp32 accumulation
lives in the 6144-length reduction chain, not in the expert count).
…ch overlapped with CPU rows Replaces the LRU-slot-tied tier (which uploaded on the decode critical path, churned with evictions, and computed ALL routed experts on the GPU while 12 cores idled — measured slower than CPU-only experts). - vk_registry_fill(): at startup, upload the top-COLI_VK_EXPERTS experts by persistent usage history into a (layer,eid) registry, decoupled from cache slots — stable residency like the HIP VRAM tier (RX 9070: 256 experts, 4.83 GB, 1.0s; pinned RAM slots feed uploads directly, the rest stream through one transient slot). - coli_vk_expert_group_issue/_take: the group gets its own command buffer + fence and splits into submit-and-return / join, so moe() ISSUES the GPU batch, computes the CPU share concurrently, drains the pipe + loads the GPU-side slots (cache invariants unchanged), then takes and accumulates. Sync coli_vk_expert_group (shared expert, harness) wraps the same path. - vk_active now keys on a non-empty registry; empty (no history/budget 0) falls back to the normal parallel CPU loop. Harness: issue/take validated bit-identical to the sync path; full PASS. Engine: 'Paris.' correct with the tier active.
…etch entirely The resolve phase now classifies registry-resident experts FIRST (before pin/LRU): they take no cache slot, dispatch no load, and skip the LRU recency bump — so they age out of RAM and free capacity for the CPU-served experts, the same effect CUDA_RELEASE_HOST gives the HIP tier. The pilot worker, cross-layer lookahead, and next-block readahead treat registry residency like RAM residency (decode only; prefill still loads normally since the tier serves S<=4). Counted as a new 'vk' bucket in the hit-rate split. The group block computes GPU entries straight from the registry (no ESlot), attributes its CPU share to t_ecpu/rows and only the take-wait to t_egpu, and on device-loss falls back via a transient ws slot load.
…ath was inverted Without a RAM pin every candidate loads transiently, and the inverted check skipped every SUCCESSFUL load (pin-fed fills masked it: the pins covered the whole top-256). PIN_GB=0 + COLI_VK_EXPERTS=256 now fills from disk in ~1.6s. Also: bail after 64 consecutive load failures, diagnostics on the first few.
256-384 measured flat within noise (1.73-1.78 tok/s); 320 is the best median with ~2.3 GB VRAM headroom left for the long-context KV mirror.
Upstream PR JustVugg#399 (KV8 fp8 / TQ4 quantized latent KV) leaves Lc/Rc NULL when a quantized tier is active — the mirror sync would deref NULL. Guard on the arrays themselves (env-independent), falling back to the CPU attention path; an fp8-aware VK mirror (upload Lc8+scale, decode e4m3 in the shader, 4x less mirror traffic) is the follow-up once JustVugg#399 merges.
Author
|
Rebased onto current dev — the split turned out painless: git tracked |
Author
|
pp512/tg128-style numbers, for comparability with the llama.cpp benchmark conventions (same box and config as the tables above; dual-SSD mirror active for both backends; fresh process per run; two runs each):
Notes: prefill here is expert-streaming-bound (a 512-token batch-union touches nearly every expert of the 744B model), so pp measures the whole engine + disks more than the GPU. The accounting slightly favors HIP — Vulkan uploads its dense weights during the first prefill while the CUDA/HIP tier uploads at load, before the prefill timer — and Vulkan leads regardless (+13% pp, +4–5% tg). These standard workloads also reproduce much more tightly than free-running generations (≤2% spread), so we'll use them for future comparisons. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
An opt-in Vulkan compute backend (
make glm VK=1, runtimeCOLI_VULKAN=1) that runs the GLM decode compute path on any GPU with a Vulkan 1.2 driver — no CUDA, no ROCm. Three tiers, each independently switchable and each falling back to the CPU path on any failure:COLI_VK_EXPERTS, default 320): the top-N routed experts by.coli_usageheat are uploaded once at startup and served from VRAM with no RAM slot, no disk read, no prefetch — issued as one fused async batch per layer (gate+up+silu→down, hidden on-device) that overlaps with the CPU computing the remaining experts. Appears as a newvkbucket in the hit-rate line.COLI_VK_DENSE=1): q_a+kv_a fused into one submit, q_b, the o-projection, and the shared expert as a single fused expert-group submit.COLI_VK_ATTN=1): one dispatch per layer covering absorbed query → scores → softmax → weighted latent → value rows, fused with the o-projection so the context vector never leaves the GPU. The latent/rope KV lives in a persistent per-layer device mirror appended at ~2.3 KB/token/layer, with the same invalidation points as the CUDA KV shadow (row rewrite, rebind, resize).This is a continuation of @kryptt's #84 (the original backend skeleton, tensor-resident
coli_vk_matmul, and the int4 nibble-8 shader decode are his design — thank you!). It addresses the two things asked for there: it is rebased onto currentdev, and it ships a CPU-reference exactness harness covering every primitive.Why
libvulkan+ a Vulkan 1.2 ICD with subgroup arithmetic, andglslcat build time.Correctness
gcc -O3 -DVK_TEST backend_vulkan.c -o test_vk -lvulkan -lm && ./test_vk shaders/qmatmul.spv): CPU-reference exactness for every primitive — int4/int8 GEMV across shapes including the long-row o-projection (I=16384), fused gate+up, the full expert group (sync and async issue/take, verified bit-identical to each other), the q_a+kv_a pair, and the absorb attention core including causal S=2,kv_startwindows, int8, and long-context (T=2000) cases. Typical maxrel 1e-5…2e-3 (fp32 reduction order on 6144-long dots).#ifdef COLI_VULKAN;make glmwithoutVK=1compiles the same code as before (plus three always-false locals).make test-cfully green on the branch.Measured performance
Hardware/software: AMD RX 9070 16 GB (RDNA4/gfx1201, PCIe 4.0 x16 effective), Ryzen 9 3900X (12C/24T Zen2,
OMP_NUM_THREADS=12, threads pinned), 64 GB DDR4-3200 dual-rank, GLM-5.2 744B int4 streamed from NVMe (~19 MB/expert). Mesa/RADV 26.1.4, Vulkan 1.4 ICD; ROCm 7.2.4 + HIP for the comparison backend. Model: 78 layers, 256 experts/layer, kv_lora 512, MLA/DSA.Primitive microbenchmarks (per-call including readback, K-expert int4 MLP 6144→2048→6144):
End-to-end GLM-5.2 decode (same box, same engine settings, expert-routing top-p 0.7, drafts off, interleaved/batched A/B runs; our fork additionally carries an unrelated dual-SSD-mirror patch on the same base — active for both backends in every comparison, so the deltas are backend-attributable):
Run-to-run variance on this box is ±0.1 tok/s; the 64-token advantage (~+13%) narrows to ~+5% at long form as the attention window and expert-I/O share grow for both backends.
The two memory-type rules that made it work (documented in
docs/vulkan.md, learned the hard way):HOST_CACHED— reading write-combined ReBAR VRAM from the CPU runs at ~40 MB/s and was a hidden 5.6× per-call penalty.HOST_VISIBLE|DEVICE_LOCAL, so uploads are plain memcpys and the GPU reads at VRAM speed.Integration surface
c/backend_vulkan.c(~1.3k lines, plain C99 + libvulkan),c/backend_vulkan.h,c/shaders/qmatmul.comp,c/shaders/qmatmul_gate_up.comp,c/shaders/attention_absorb.comp,docs/vulkan.md.c/glm.c(hooks under#ifdef COLI_VULKAN),c/Makefile(VK=1block +glslcrule),docs/ENVIRONMENT.md,README.md.COLI_VULKAN,COLI_VK_SHADERS,COLI_VK_EXPERTS,COLI_VK_DENSE,COLI_VK_ATTN.One operational note worth flagging beyond this PR: the engine's OMP self-tune (
OMP_WAIT_POLICY=active+ spin) is skipped underCOLI_CUDA/COLI_METALbut not under other configurations. With 12 pinned spinning threads the async I/O pool starves (measured 28 → 5 GB/s CPU expert bandwidth, ~2.8× end-to-end).docs/vulkan.mdtells Vulkan users to setCOLI_NO_OMP_TUNE=1; extending the self-tune skip to any GPU backend (or toPIPE=1runs) may be worth a separate look.Limits / future work
S≤4; prefill keeps the CPU/batched paths (dense projections do run on VK at prefill, which is why prefill got faster).KV8=1(already implemented and validated againstkv_fp8.hon our side — the absorb shader decodes e4m3 in place, 4× less mirror traffic; this PR just cleanly falls back to CPU attention when the f32 cache is absent), cooperative-matrix prefill kernels, and a fully resident-layer pipeline.How to run