metal: implement DSv4 Lightning Indexer - #25893
Conversation
- Implement GGML_OP_LIGHTNING_INDEXER for 128-dimensional, 64-head inputs with F32 queries and weights plus F16 keys and masks. - Add tiled and tail kernels and test KV lengths around 8- and 64-element boundaries. llama-bench (--mmap 1, -fa 1, -p 512, -n 128; d=0/10k/20k/30k): Before: - pp512: 153.73 ± 0.87 t/s - tg128: 8.91 ± 0.04 t/s - pp512 @ d10000: 73.90 ± 0.39 t/s - tg128 @ d10000: 8.66 ± 0.03 t/s - pp512 @ d20000: 45.83 ± 0.18 t/s - tg128 @ d20000: 8.26 ± 0.03 t/s - pp512 @ d30000: 33.40 ± 0.21 t/s - tg128 @ d30000: 7.94 ± 0.01 t/s After: - pp512: 155.19 ± 0.91 t/s - tg128: 8.95 ± 0.04 t/s - pp512 @ d10000: 86.95 ± 0.69 t/s - tg128 @ d10000: 9.00 ± 0.05 t/s - pp512 @ d20000: 62.01 ± 0.45 t/s - tg128 @ d20000: 8.68 ± 0.04 t/s - pp512 @ d30000: 49.18 ± 0.33 t/s - tg128 @ d30000: 8.60 ± 0.02 t/s Assisted-by: Codex
github-actions
Bot
added
testing
|
@ggml-org/ggml-metal can someone review? |
|
Since creating this, I've pushed GPT 5.6 Sol to add more optimizations to my dsv4-improvements branch. These are the latest numbers of DSv4 on my M1 ultra: In my tests, it is working better than antirez DS4 llama.cpp fork. Besides the extra optimizations, using Besides the metal lightning indexer commit in this PR, there are about 10 commits. I can push them all together in a single PR, create stacked PRs or individual PRs (whatever is best for maintainers to review). |
|
@am17an thanks I will try re-converting it later. if anyone wants to try the new 0731 gguf, I'm currently uploading IQ3_XXS with the 0731 template already embedded (and without MTP) here: https://huggingface.co/tarruda/DeepSeek-V4-Flash-0731-GGUF. IQ3_XXS seems to be the sweet spot for 128G |
|
@tarruda Thanks for the contribution — this is a nice improvement. A couple of comments.
I'd rather handle the tail the same way we already do in I'd also rather not make this FP16-only. DSv4 is intended for very long contexts, where a quantized K cache is likely to be the configuration people will want to run. The only thing preventing both is simdgroup_load(mk[i], (device const half *) k_base + 8*i, ...)
That's already what the flash-attention quantized-K path does. The dequantization helpers already exist with a uniform interface, and kernel instantiation already follows the Once K is staged, the tail case becomes what CUDA already does in // staging
if (i_kv + r < args.n_kv) { /* load + dequantize K row */ } else { /* zero-fill */ }
// store
if (ik < args.n_kv) out[ik] = score + float(mask[ik]);Zero-filling K naturally gives That lets us drop the tail kernel, support quantized K as well, and simplify |
- Stage and dequantize K in F16 threadgroup memory before simdgroup matrix loads. - Zero-fill partial tiles and guard stores so all KV segments use the same numerical path. - Support F32, F16, BF16, Q4_0, Q4_1, Q5_0, Q5_1, and Q8_0 K caches. llama-bench (--mmap 1, -fa on, -p 512, -n 128; d=0/10k/20k): - pp512: 160.38 +/- 1.01 t/s - tg128: 9.08 +/- 0.03 t/s - pp512 @ d10000: 88.37 +/- 0.46 t/s - tg128 @ d10000: 9.07 +/- 0.04 t/s - pp512 @ d20000: 62.53 +/- 0.46 t/s - tg128 @ d20000: 8.84 +/- 0.03 t/s Assisted-by: Codex
|
@forforever73 merged master branch and pushed a new commit. In the interest of full disclosure, I'd like to reinforce that I don't know what I'm doing and all I did was forward your comment to GPT 5.6 Sol. I'm happy to keep using my codex subscription on this, but my blessing goes to you or anyone else that wants to create a new PR with this code cleaned up, improved or rewritten. This also applies to my dsv4-improvements branch. I'm happy as long as we end up getting good performance improvements on DSv4 inference. |
|
@tarruda Thanks for being so open about it :) At the end of the day we're all after the same thing: getting the best possible DSv4 performance. From a quick look, 5.6 Sol seems to have captured the direction I was suggesting. I'll probably just make a few changes directly on top of your branch to keep things moving. |
|
BTW, on top of this the HC ops probably seem to provide the most benefit for TG. Those should be a quick merge |
Most commit messages in my branch contain llama-bench measurements for the change. And yes, HC greatly improve token generation: 6ed571d Another great improvement is in the sparse attention commits which follow, making prompt processing more stable with longer prefills: 551afb9 |
|
@forforever73 Yes, feel free to move forward - I haven't gotten around to setup DSv4 in my environment yet, so any help with this is appreciated. |
|
The first one is similar to #25917, though I am not sure if that will be merged soon. The second commit also looks like a bit along the same lines. |
|
@ggerganov no, I was going to open PRs in the same order of the commits as needed, but feel free to take anything out of that branch and create PRs yourself |
Ok, I'll fast-track that specific commit: #26459 |
|
Thanks @ggerganov, let me know if I can do anything else to assist. |
Ok, let's wait for @forforever73 to update this branch and we can continue with the rest of the changes from your branch after merging the lightning indexer. |
Assisted-by: Codex # Conflicts: # ggml/src/ggml-metal/ggml-metal-device.cpp # ggml/src/ggml-metal/ggml-metal-device.h # ggml/src/ggml-metal/ggml-metal-device.m # ggml/src/ggml-metal/ggml-metal-impl.h # ggml/src/ggml-metal/ggml-metal-ops.cpp # ggml/src/ggml-metal/ggml-metal-ops.h # ggml/src/ggml-metal/ggml-metal.metal
Ok just pushed a merge commit with conflicts fixed for @forforever73 |
|
Test on my m4 max, it takes a long time
|
|
@forforever73 I think you started the test before the hyper-connections kernel was merged back into this PR (thus the 10 tps generation). The pp2048 will greatly improve at longer context once sparse attention commits are added. |
|
@tarruda Yes, exactly. So I think the current results are expected. |
forforever73
added
the
merge ready
|
Some quick numbers on M2 Ultra with https://huggingface.co/ggml-org/DeepSeek-V4-Flash-0731-GGUF: |
|
@ggerganov these numbers are more than double the pp2048 I get with an M1 ultra. Can you share the llama-bench command you used? Curious if I missed some flags or if the M2 is that much better than the M1. |
|
I'm using your command from earlier: sudo sysctl iogpu.wired_limit_mb=180000
make -j && bin/llama-bench -hf ggml-org/DeepSeek-V4-Flash-0731-GGUF -fa 1 -d 0,10000,20000,30000,40000,50000,60000,70000,80000,90000,100000,150000,200000,250000,300000,350000,400000,500000 -n 128 -p 2048 -b 2048 -ub 2048Do you have this patch: #26465 |
|
Yea I did. I guess it is just that the M1 GPU is significantly weaker than the M2. BTW, I think the next best picks from my branch are the sparse attention ones which keep prompt processing more stable as context grows. Do you want me to open a new PR or want to pick the changes yourself? |
Yes, go ahead and open - I'll take a look. |
|
Hi @tarruda, this PR will cause the github action M1 Virtual Machine raise a error: 
|
|
@JamePeng do you know how I can run this test locally? Seems like it is causing an issue on unrelated model |
|
It seems the paravirtual device does not allow to allocate I think you need to reserve the threadgroup memory in host-side using |
|
@ggerganov I'm going to open a PR, but how can we validate the fix before merging? Was this not caught in the PR CI? |
|
Our M1 virtual runners announce that simdgroup matrix multiplication is not supported, so the CI filters out those tests as unsupported: https://github.com/ggml-org/llama.cpp/actions/runs/30784775242/job/91596148481#step:5:20121 Not sure how @JamePeng runs this. |
Simply loading the Qwen2.5-0.5B model and conducting some grammar tests. |
|
I submitted a fix #26646 , and it worked. |
* metal: implement F16 Lightning Indexer - Implement GGML_OP_LIGHTNING_INDEXER for 128-dimensional, 64-head inputs with F32 queries and weights plus F16 keys and masks. - Add tiled and tail kernels and test KV lengths around 8- and 64-element boundaries. llama-bench (--mmap 1, -fa 1, -p 512, -n 128; d=0/10k/20k/30k): Before: - pp512: 153.73 ± 0.87 t/s - tg128: 8.91 ± 0.04 t/s - pp512 @ d10000: 73.90 ± 0.39 t/s - tg128 @ d10000: 8.66 ± 0.03 t/s - pp512 @ d20000: 45.83 ± 0.18 t/s - tg128 @ d20000: 8.26 ± 0.03 t/s - pp512 @ d30000: 33.40 ± 0.21 t/s - tg128 @ d30000: 7.94 ± 0.01 t/s After: - pp512: 155.19 ± 0.91 t/s - tg128: 8.95 ± 0.04 t/s - pp512 @ d10000: 86.95 ± 0.69 t/s - tg128 @ d10000: 9.00 ± 0.05 t/s - pp512 @ d20000: 62.01 ± 0.45 t/s - tg128 @ d20000: 8.68 ± 0.04 t/s - pp512 @ d30000: 49.18 ± 0.33 t/s - tg128 @ d30000: 8.60 ± 0.02 t/s Assisted-by: Codex * metal: stage Lightning Indexer K tiles - Stage and dequantize K in F16 threadgroup memory before simdgroup matrix loads. - Zero-fill partial tiles and guard stores so all KV segments use the same numerical path. - Support F32, F16, BF16, Q4_0, Q4_1, Q5_0, Q5_1, and Q8_0 K caches. llama-bench (--mmap 1, -fa on, -p 512, -n 128; d=0/10k/20k): - pp512: 160.38 +/- 1.01 t/s - tg128: 9.08 +/- 0.03 t/s - pp512 @ d10000: 88.37 +/- 0.46 t/s - tg128 @ d10000: 9.07 +/- 0.04 t/s - pp512 @ d20000: 62.53 +/- 0.46 t/s - tg128 @ d20000: 8.84 +/- 0.03 t/s Assisted-by: Codex * dedup Lightning Indexer constants, fix flaky test * cont : fix whitespace --------- Co-authored-by: forforever73 <690105611@qq.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* metal: implement F16 Lightning Indexer - Implement GGML_OP_LIGHTNING_INDEXER for 128-dimensional, 64-head inputs with F32 queries and weights plus F16 keys and masks. - Add tiled and tail kernels and test KV lengths around 8- and 64-element boundaries. llama-bench (--mmap 1, -fa 1, -p 512, -n 128; d=0/10k/20k/30k): Before: - pp512: 153.73 ± 0.87 t/s - tg128: 8.91 ± 0.04 t/s - pp512 @ d10000: 73.90 ± 0.39 t/s - tg128 @ d10000: 8.66 ± 0.03 t/s - pp512 @ d20000: 45.83 ± 0.18 t/s - tg128 @ d20000: 8.26 ± 0.03 t/s - pp512 @ d30000: 33.40 ± 0.21 t/s - tg128 @ d30000: 7.94 ± 0.01 t/s After: - pp512: 155.19 ± 0.91 t/s - tg128: 8.95 ± 0.04 t/s - pp512 @ d10000: 86.95 ± 0.69 t/s - tg128 @ d10000: 9.00 ± 0.05 t/s - pp512 @ d20000: 62.01 ± 0.45 t/s - tg128 @ d20000: 8.68 ± 0.04 t/s - pp512 @ d30000: 49.18 ± 0.33 t/s - tg128 @ d30000: 8.60 ± 0.02 t/s Assisted-by: Codex * metal: stage Lightning Indexer K tiles - Stage and dequantize K in F16 threadgroup memory before simdgroup matrix loads. - Zero-fill partial tiles and guard stores so all KV segments use the same numerical path. - Support F32, F16, BF16, Q4_0, Q4_1, Q5_0, Q5_1, and Q8_0 K caches. llama-bench (--mmap 1, -fa on, -p 512, -n 128; d=0/10k/20k): - pp512: 160.38 +/- 1.01 t/s - tg128: 9.08 +/- 0.03 t/s - pp512 @ d10000: 88.37 +/- 0.46 t/s - tg128 @ d10000: 9.07 +/- 0.04 t/s - pp512 @ d20000: 62.53 +/- 0.46 t/s - tg128 @ d20000: 8.84 +/- 0.03 t/s Assisted-by: Codex * dedup Lightning Indexer constants, fix flaky test * cont : fix whitespace --------- Co-authored-by: forforever73 <690105611@qq.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
5 participants
Overview
Metal backend for GGML_OP_LIGHTNING_INDEXER
Additional information
This is the first of a series of changes that target improving DSv4 performance on Apple Silicon. I extracted this commit from my dsv4-improvements branch, which includes metal backends for hyper-connections and a few other improvements.
On master branch, these are some performance numbers of DSv4 on my M1 Ultra (llama-bench --mmap 1, -fa 1, -p 512, -n 128; d=0/10k/20k/30k, benchmarking this quant):
With this commit, things improve slightly:
With all commits from my branch, this is what it looks like:
Requirements