Dsv4 metal - #3
Conversation
…eepSeek V3.2/V4 lightning indexer
…gml_top_k() implementation to reduce temporary buffers memory usage
Implements the DeepSeek V4 lightning indexer on Metal GPU. Follows the CUDA vec kernel approach with 8 SIMD groups per threadgroup, each processing 8 KV vectors using simd_sum for per-head dot product reduction. Supports F32, F16, BF16 and quantized K types (Q4_0, Q4_1, Q5_0, Q5_1, Q8_0). Assisted-by: DeepSeek V4 Pro
Implements GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, and GGML_OP_DSV4_HC_POST on Metal GPU. HC_PRE performs weighted sum over hc slices, HC_COMB computes Sinkhorn-normalized combination matrices, and HC_POST blends input with residuals using the combination weights. All kernels operate on F32 tensors. Assisted-by: DeepSeek V4 Pro
The Metal lightning indexer assigned a scalar float expression to the float4 threadgroup q tile and derived the address from the packed embedding width. That ignored the q head stride and loaded the wrong q vector data, corrupting indexer scores for DeepSeek V4 on Metal. Load q tiles as strided float4 values, matching the CUDA path, and use the provided source and destination strides in HC_PRE instead of assuming contiguous row layout. Add a LIGHTNING_INDEXER backend-op case so the DSV4-shaped F32 path is checked against CPU. Assisted-by: Codex
|
@tarruda I'm not interested in maintaining my own separate fork, so I've been submitting all my changes as PRs to mainline llama.cpp, now it's time for lightning indexer: ggml-org#24231 and ggml-org#25545 It currently lacks Metal implementation, I think it would be useful. Maybe extract it as a separate PR and submit as PR to ggm-org/llama.cpp? Note however that lightning indexer API changed slightly - now it no longer does the scaling (easy to fix) and a f16 mask parameter was added to it. There are lightning indexer tests in the CPU PR ggml-org#24231 so you can also check correctness and measure performance. Leaving this open for other people until mainline supports it. Good luck! |
dd9b87c to
c74cb24
Compare
|
Closing as I have created a better metal branch on top of llama.cpp mainline (which now has lightning indexer and hyperconnections). I will see if I can extract the changes and submit individually to upstream |
…rg#24908) * Add preliminary MiniMax-M3 support Text-only port that re-uses existing components: MiniMax-M2 style GQA with per-head QK-norm and partial rotary, DeepSeek-V3 style leading-dense and routed/shared experts, and swigluoai activation. Sparse attention is not yet supported (dense fallback); vision tower and MTP heads are dropped. * MiniMax-M3 vision tower (mmproj + clip graph) * Delete m3_vision_ref.py * Update clip.cpp * MSA * Update constants.py * Update minimax.py * Cache creation. Working withotu flash attention * Added flash attention for sparse layers * Decomposed slow cpu OP into GPU + CPU ops. Massive speedup over long ctx * Rewrote indexer op to be cuda native. Modified flash attention to match per group block picking * Implement sparse attention calc out of stock ops. * Fix a cache allocation and cont issue * Fixed -fa auto crash, flagged debug spots * Delete vocab.json * Delete model.safetensors.index.json * Delete generation_config.json * Delete Minimax directory * Handled multi stream case to fall back on Dense Attention * Development scaffolding cleanup. No functional change to the decode or 4-way paths. Full debug harness remains at <8136a9c68ed7a5eb009aa67bba3fda8062f4648f> for reproducing the selection-parity validation. * Remove redundant comment from minimax-m3.cpp * Changed 3 Gelu Ops for vision into Gelu_erf ops * Assert that n_kv is multiple of 128 * Rename MSA index tensors to indexer convention Note: All GGUFs generated before this change will need to be regenerated. * Fix incorrect Assert * Review driven changes (#3) * Remove comment from conversion minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespaces from constants.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Tighten comment in minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * inherit MiniMax-M3 from MiniMax-M2 * drop dead text_config fallbacks * Add indexer writer methods * Reuse LLM_FFN_SWIGLU_OAI_MOE * Remove duplicate indexer setters, add only block_size/local_blocks, follow value naming convention * Fix conversion error /gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/tensor_mapping.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-kv-cache.cpp Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove Whitespace in Update src/llama-model.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-hparams.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * remove multimodal code upon maintainer request. Will be made as a separate PR * Whitespace clean in tensor_mapping.py * Log cache size on launch, block ctx shift, support prompt caching Log indexer cache size on launch Disallow ctx shift Support prompt caching * Update minimax-m3.cpp * Optimize implementation, add multi stream support. Fully rewrote minimax-m3.cpp for speed and buffer size gains: Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3] Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill Decode: ~25 nodes/layer vs ~50, no per-group concats/conts Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token) In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq Measured effect on expert offload bound setup: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support. * set default cache type to F32 * Fix potential DSA double indexer cache allocation bug, only allocate in-cache k_idx for archs that opt in * remove F16 downcasts in MSA attention, force F32 indexer score accum * Add Minimax eos to llama vocab * Guard edge case where idx cache can become stale after a tail trim * Update llama-kv-cache.h * Update llama-kv-cache.cpp * Update llama-kv-cache.cpp * Update llama-kv-cache.h * Update llama-kv-cache.cpp * Review driven changes * style fix * indexer hparams are required * fix tests * fix lint --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
* Add preliminary MiniMax-M3 support Text-only port that re-uses existing components: MiniMax-M2 style GQA with per-head QK-norm and partial rotary, DeepSeek-V3 style leading-dense and routed/shared experts, and swigluoai activation. Sparse attention is not yet supported (dense fallback); vision tower and MTP heads are dropped. * MiniMax-M3 vision tower (mmproj + clip graph) * Delete m3_vision_ref.py * Update clip.cpp * MSA * Update constants.py * Update minimax.py * Cache creation. Working withotu flash attention * Added flash attention for sparse layers * Decomposed slow cpu OP into GPU + CPU ops. Massive speedup over long ctx * Rewrote indexer op to be cuda native. Modified flash attention to match per group block picking * Implement sparse attention calc out of stock ops. * Fix a cache allocation and cont issue * Fixed -fa auto crash, flagged debug spots * Delete vocab.json * Delete model.safetensors.index.json * Delete generation_config.json * Delete Minimax directory * Handled multi stream case to fall back on Dense Attention * Development scaffolding cleanup. No functional change to the decode or 4-way paths. Full debug harness remains at <8136a9c68ed7a5eb009aa67bba3fda8062f4648f> for reproducing the selection-parity validation. * Remove redundant comment from minimax-m3.cpp * Changed 3 Gelu Ops for vision into Gelu_erf ops * Assert that n_kv is multiple of 128 * Rename MSA index tensors to indexer convention Note: All GGUFs generated before this change will need to be regenerated. * Fix incorrect Assert * Review driven changes (#3) * Remove comment from conversion minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespaces from constants.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Tighten comment in minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * inherit MiniMax-M3 from MiniMax-M2 * drop dead text_config fallbacks * Add indexer writer methods * Reuse LLM_FFN_SWIGLU_OAI_MOE * Remove duplicate indexer setters, add only block_size/local_blocks, follow value naming convention * Fix conversion error /gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/gguf_writer.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update gguf-py/gguf/tensor_mapping.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update conversion/minimax.py Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-kv-cache.cpp Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove Whitespace in Update src/llama-model.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Remove whitespace in src/llama-hparams.h Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> * Update minimax_m3.cpp Rewrite code comment based on feedback and to better reflect the actual architecture, and reuse existing build_vit * Rename minimax_m3.cpp to minimax-m3.cpp * Update CMakeLists.txt * Remove debug code from clip.cpp * Update clip.cpp * Update comments in tools/mtmd/models/minimax-m3.cpp * Permute Q/K at conversion, drop precomputed sin/cos * Log cache size on launch, block ctx shift, support prompt caching Log indexer cache size on launch Disallow ctx shift Support prompt caching * Update minimax-m3.cpp * Optimize implementation, add multi stream support. Fully rewrote minimax-m3.cpp for speed and buffer size gains: Unified the 4-way + decode, 1 FA call per layer instead of 4, with the groups mapped onto ne[3] Custom CPU op now emits block-level mask, expanded on GPU, which causes CPU to GPU transfer to shrinks at prefill Decode: ~25 nodes/layer vs ~50, no per-group concats/conts Unified selection semantics, so both regimes rank bs + local bias (position-anchored local force), which means prefill/decode can no longer disagree on selection can_reuse on the MSA bias input. Graph reuse at decode restored (was rebuilding the full graph every token) In-place mask adds, shrinking compute buffer ~6.8 to ~4.2 GiB at ub2048/62k Multi-stream: MSA now runs with -np N when kv_unified=false. Decode stays batched across streams (still 1 FA call), prefill loops per stream. dense fallback only for --kv-unified + multi-seq Measured effect on expert offload bound setup: decode 6.2(4WAY)–7.15(MSA_decode) -> 7.7~7.8 t/s, flat from 5k to 60k+. prefill around 10% faster. buffer about 20% smaller, multi-user support. * set default cache type to F32 * Fix potential DSA double indexer cache allocation bug, only allocate in-cache k_idx for archs that opt in * remove F16 downcasts in MSA attention, force F32 indexer score accum * Add Minimax eos to llama vocab * Guard edge case where idx cache can become stale after a tail trim * Update llama-kv-cache.h * Update llama-kv-cache.cpp * Update llama-kv-cache.cpp * Update llama-kv-cache.h * Change resize Pad to none, resize alg to Bicubic Pillow * Review driven changes * Update llama-kv-cache.cpp * rm unrotated pos_t * fused rope w + pad * rename merge --> merger for consistency * add review skill for mtmd * graph should use hparams n_merge * fix lint --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
This was mostly vibe coded by Deepseek V4 Pro, with the last commit being Codex fixing a lightning op metal bug introduced in the previous commits (plus a lightning op backend test).
On my M1 Ultra, this increases DSV4-flash token generation speed to ~20tps from ~6tps. I've been running this with this IQ3_XXS quant.
From my testing this looks like it is working well. Unlike the llama.cpp upstream implementation, it doesn't seem to have this bug, or at least I couldn't reproduce after a pi session that used more than 100k tokens in context.
IDK if you have restrictions against AI coded changes or if you are interested in merging. But since the dsv4 branch doesn't seem to be meant for a llama.cpp PR I thought it would be good to have it centralize implementation for more backends, which can be used as reference for future llama.cpp PRs.