metal : per-device tuned (Q, NE) for flash-attn vec - #25750
Conversation
|
@forforever73 I've updated the |
50b90d0 to
4d5dffa
Compare
|
Done, ready for review now |
ggerganov
left a comment
There was a problem hiding this comment.
The changes to test-backend-ops.cpp are very heavy. Should be a separate binary?
|
@ggerganov Yes, that's the plan. I'm going to add a separate tool for this. |
Ah, ok. I just saw the follow-ups in the description. |
| if (strcmp(name, "ggml_backend_metal_set_fa_vec_override") == 0) { | ||
| return (void *)ggml_backend_metal_set_fa_vec_override; | ||
| } | ||
| if (strcmp(name, "ggml_backend_metal_clear_fa_vec_override") == 0) { | ||
| return (void *)ggml_backend_metal_clear_fa_vec_override; | ||
| } | ||
| if (strcmp(name, "ggml_backend_metal_fa_vec_ne11_bucket") == 0) { | ||
| return (void *)ggml_backend_metal_fa_vec_ne11_bucket; | ||
| } | ||
| if (strcmp(name, "ggml_backend_metal_fa_vec_ne01_bucket") == 0) { | ||
| return (void *)ggml_backend_metal_fa_vec_ne01_bucket; | ||
| } | ||
| if (strcmp(name, "ggml_backend_metal_fa_vec_baseline_ne") == 0) { | ||
| return (void *)ggml_backend_metal_fa_vec_baseline_ne; | ||
| } |
There was a problem hiding this comment.
The names of these procs should have the prefix ggml_backend_metal_tuning_.
For example: ggml_backend_metal_tuning_set_fa_vec_override
| // test/tune-only override; when set, fa_vec_pick returns it directly. | ||
| void fa_vec_set_override(fa_vec_cfg_t cfg); | ||
| void fa_vec_clear_override(); | ||
| bool fa_vec_override_active(); |
There was a problem hiding this comment.
fa_vec_override_active seems unused
| struct fa_vec_key_t { | ||
| int8_t device_id; | ||
| int8_t dtype; | ||
| int16_t dk; | ||
| int16_t dv; | ||
| int8_t ne11_b; | ||
| int8_t ne01_b; | ||
| }; |
There was a problem hiding this comment.
At some point we will also try to add auto-tuning for other ops, such as MUL_MAT and MUL_MAT_ID. Not sure if it's worth to generalize the tuning key to work with all ops. Alternatively, we can add individual key structs for each op. Just marking this for consideration later on.
There was a problem hiding this comment.
Inspired by #25377 and based on my current experiments, adding auto-tuning for mul_mm also shows improvements. The tuning parameters differ somewhat from those for fa_vec, though. Once this PR is merged, I'll take a step back and evaluate whether it makes sense to generalize the tuning key or keep separate key structs for different ops.
ggerganov
left a comment
There was a problem hiding this comment.
You should be able to merge this, since the base branch is not master.
…ed fa_vec_override_active
Overview
cont #23566.
This PR adds a compile-time lookup table that picks a better
(Q, NE)per(device, dtype, head size, KV depth, batch width), following the generic-tuning / per-device-config direction from the #23114 discussion.The table covers 6 KV dtypes: f16 plus the five quantized caches
q4_0/q4_1/q5_0/q5_1/q8_0. Zero regression when the table doesn't hit: the baseline path stays byte-identical to today's backend in host_name, pipeline cache key, and numerics — only a table hit selects a new_q{Q}_ne{NE}kernel.How it works
test-backend-ops tune --tune-perfsweeps the grid offline and prints a pasteable C table, so any Apple-silicon owner can reproduce it.fa_vec_pickmatches the device-exact SKU first, then the family's representative SKU — same-family parts benefit without their own sweep.The table selects only Q and NE. nsg/nwg stay on the existing adaptive heuristic.The tuning table
The key is 2D per dtype —
(ne11_bucket, ne01_bucket):{1024, 4096, 16384}; the baseline→Q>1 crossover is head-size dependent. Short KV (ne11 < 1024) is left to baseline: attention is a small share of the step at short context, and short-KV timings are the noisiest.{2,3,4,5}split decode (ne01==1) from batch (≥2) and point-resolve ne01 = 2, 3, 4. The per-point-optimal Q isn't monotonic in the batch: q5's heavy dequant favors Q4 at ne01=3–4 while other dtypes stay Q1/Q2. Since real draft-n-max is typically ≤3 (ne01 ≤ 4), ne01 ≥ 5 collapses into a single min-max-regret backstop bucket.The table is stored group-split to stay small: within one
(dtype, dk, dv)the winning cfg is usually constant across ne11 buckets, so the generator emits one ne11-collapsed default row per ne01 domain, plus exact-bucket exception rows only where a bucket would lose too much on that default.Two thresholds gate this. θ = 1.05 (target selection): a bucket is tuned only if its target beats baseline by ≥5% in aggregate over its sampled points, else it keeps baseline. τ = 5% (compression): a bucket folds into the domain default only if its pointwise worst regret vs its own target stays within 5%, else it becomes an exception.
Kernel micro-benchmarks (M4 Max, Apple9)
test-backend-ops tune --tune-perfsweeps the(Q, NE)grid on a GQA spec-decode shape (nh=4, nr23=[8,1],type_K==type_V). Each candidate's time is the per-op median of 7 reps, timed in shuffled order with a periodic baseline anchor to flag thermal throttling. The full per-(dtype, dk, ne11, ne01)matrix is in the attached sweep.What
fa_vec_pickcommits at runtime for a dk=128, ne11=32768 head, per dtype (decode ne01=1, batch ne01=2–4). These are the table's actual picks, not the per-point best. Each stays within τ=5% of it (e.g. q5_0 ne01=3/4 commit Q2 where this point's fastest is Q4). ne01=1 stays Q1 and is numerically identical to baseline, so a speedup > 1 there is a pure NE retune on caches whose baseline NE was suboptimal:For the tail bucket (ne01 ≥ 5), one cfg must cover the whole range, and no single point's winner generalizes — a higher-Q cfg wins on Q-aligned ne01 but wastes padded rows on misaligned ne01. So the runtime picks the min-max-regret cfg: the one minimizing the worst per-point regret over the bucket, measured against the per-point oracle.
On a representative bucket (
q4_0, dk=dv=192, ne11=2048) the per-point winner flips: Q1,NE4 is fastest at ne01=5–7, Q2,NE2 at ne01=8/16 — no single cfg is optimal across the bucket:t/t_oracle − 1)So the runtime commits Q1,NE4 — within 3.3% of the per-point best at every ne01 (mean 0.7%). The two shortcuts blow up at the misaligned small-ne01 points: extrapolating the single ne01=8 winner (Q2,NE2) is 19.4% off at ne01=5, a uniform Q2,NE4 up to 36.5%.
In real spec-decode,
--spec-draft-n-max ≤ 3keeps ne01 ≤ 4 on the point buckets (regret ≈ 0); the tail bucket mainly backstops large draft-n / MTP.The raw
--tune-perfoutput every table above:fa_vec_sweep_m4max.txt
End-to-end (speculative decoding)
llama-server -m <target> [-md <draft>] --spec-type draft-simple|draft-mtp --spec-draft-n-max 3 -fa on -np 4, greedy. A fixed-token prompt pins the KV depth (ne11); decode tok/s is the servertimingsmedian of 3, baseline build (empty table) vs tuned build.Across the 48 comparisons (4 targets × {f16, q8_0, q5_0, q4_0} × {4k, 16k, 32k}), speedup grows with KV length, peaking at 1.323× (Qwen3-30B-A3B q5_0 @32k). The only two rows below 1.0× — 27B q4_0 @4k (0.994×) and 35BA3B q8_0 @32k (0.988×) — are draft-mtp cases: the tuned kernel (Q≥2 in verify) slightly perturbs the MTP head's logits and shifts which drafts are accepted (35BA3B q8_0 @32k: acceptance drops from 0.99 baseline to 0.72 tuned). Dividing that out, the kernel step-rate is 1.027× and 1.243×, and grows with depth like the rest.
Full e2e matrix — 4 targets × {f16, q8_0, q5_0, q4_0} × {4k, 16k, 32k} (48 comparisons)
Compilation cost
Adding a new device
The table is collected per-device with the same offline sweep that produced the M4 Max rows, so any Apple-silicon owner can extend it.
Contributor workflow
cmake --build <build> --target test-backend-ops -j./build/bin/test-backend-ops tune --tune-perf -b MTL0 > fa_vec_sweep.txt 2>/dev/nullThe part that gets folded into the table is the
fa_vec_tuned_tableblock at the end — one// ---- <dtype>: N rows ----section per dtype. It's finished output (min-max-regret target, θ=1.05 gate, short-KV drop, and τ=5% compression already applied); the attachedfa_vec_sweep_m4max.txtis exactly this format../build/bin/test-backend-ops tune -b MTL0(drift guard + per-(dtype, Q, NE) matrix vs CPU-ref)./build/bin/test-backend-ops test -o FLASH_ATTN_EXT -b MTL0fa_vec_sweep.txtin the discussion thread.Validation output (this machine,
build_embed)tune -b MTL0covers f16 + all 5 quantized caches × every(Q, NE)× padded rows × sinks × nsg(1/2/4); the quantized Q>1 dequant-once / cross-Q reuse is correct at all precisions.Once the discussion thread has collected sweeps from other devices, I'll fold each one's
fa_vec_tuned_tableblock into the table — registering the SKU enum inggml-metal-device.{h,m}where a device isn't listed yet (the M1–M5 families already are, so this is rare).follow-ups
# WARN); a cooldown mechanism (run a batch, pause until the GPU cools, resume) would keep long sweeps throttle-free instead of needing a re-run.tests/test-backend-ops.cppinto a standalonetools/metal-kernel-tuning(FA-vec as the first plugin); design indocs/superpowers/specs/2026-07-13-metal-kernel-tuning-tool.md.Requirements