New GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer

[fairydreaming]

Overview

This PR adds new GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer. The purpose of this OP is to reduce compute buffer sizes. Savings are enormous, for example:

No GGML_OP_LIGHTNING_INDEXER

$ ./bin/llama-batched-bench -m ../models/DeepSeek-V3.2-Q8_0.gguf -ub 2048 -npl 1 -npp 2048 -ntg 32 -fa 1 --verbose
0.00.038.860 I llama_model_loader: loaded meta data with 59 key-value pairs and 1420 tensors from ../../llama.cpp-deepseek-v32-minimal/models/DeepSeek-V3.2-Q8_0.gguf (version GGUF V3 (latest))
...
0.46.752.636 I sched_reserve:        CPU compute buffer size = 168368.12 MiB
...

With GGML_OP_LIGHTNING_INDEXER

$ ./bin/llama-batched-bench -m ../models/DeepSeek-V3.2-Q8_0.gguf -ub 2048 -npl 1 -npp 2048 -ntg 32 -fa 1 --verbose
0.00.041.225 I llama_model_loader: loaded meta data with 59 key-value pairs and 1420 tensors from ../../llama.cpp-deepseek-v32-minimal/models/DeepSeek-V3.2-Q8_0.gguf (version GGUF V3 (latest))
...
0.46.721.986 I sched_reserve:        CPU compute buffer size =  5808.12 MiB
...

Performance on CPU is unchanged (actually it's slightly faster with GGML_OP_LIGHTNING_INDEXER on my machine):

No GGML_OP_LIGHTNING_INDEXER

$ ./bin/llama-bench -m ../models/DeepSeek-V3.2-Q8_0.gguf -fa 1 -p 512 -n 32 -r 3
| model                          |       size |     params | backend    | threads |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | ------: | --: | --------------: | -------------------: |
| deepseek32 ?B Q8_0             | 678.56 GiB |   685.36 B | CPU        |      32 |   1 |           pp512 |         25.21 ± 0.01 |
| deepseek32 ?B Q8_0             | 678.56 GiB |   685.36 B | CPU        |      32 |   1 |            tg32 |          6.29 ± 0.00 |

build: 5a69c9743 (9539)

With GGML_OP_LIGHTNING_INDEXER

$ ./bin/llama-bench -m ../models/DeepSeek-V3.2-Q8_0.gguf -fa 1 -p 512 -n 32 -r 3
| model                          |       size |     params | backend    | threads |  fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | ------: | --: | --------------: | -------------------: |
| deepseek32 ?B Q8_0             | 678.56 GiB |   685.36 B | CPU        |      32 |   1 |           pp512 |         25.51 ± 0.01 |
| deepseek32 ?B Q8_0             | 678.56 GiB |   685.36 B | CPU        |      32 |   1 |            tg32 |          6.52 ± 0.01 |

build: 5a69c9743 (9539)

Additional information

DeepSeek lightning indexer torch implementation is as follows (taken from DeepSeek V4 model.py):

weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5)
index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio])
index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2)

The problem with the naive GGML implementation of this is that the einsum matrix multiplication produces a temporary result with size proportional to the ubatch size, kv cache length and the number of indexer heads (64). When using the full context of DeepSeek V3.2 (163840) or DeepSeek V4 (1048576/4) this will take a lot of memory. By adding a specialized OP that fuses all operations we can reduce the compute buffer memory size 64 times (the number of indexer heads).

Next Steps

The code was initially taken straight from #21149, but I think there are still some decisions to be made:

• naming (use any abbreviations or not),
• whether to include scale factors in arguments or leave them outside by prescaling indexer weights,
• whether to perform dot product calculations from fp32 Q and fp32 dequantized K like I did initially or from K and Q quantized to the indexer K type (like in the GGML matrix multiplication implementation),
• consider adding mask parameter to this OP, so that we don't have to add the mask to the indexer scores (mask could be f16 then).

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: NO

[fairydreaming]

Pinging @am17an for any DeepSeek V4 related suggestions.

[am17an]

I think not materialising the intermediate tensor would be crucial to make dsv4 work because this can't be solved by op fusion. I'm thinking there should be a better way to add new ops which don't fall back to CPU in case the implementation is not present in a backend. We do this for FA and GDN, I'm wondering if we can refactor the code to make it easier to add any op like this (i.e. have a decomposed ggml fallback instead of relying on CPU impl), in that case the maintainability burden would be greatly reduced. cc @ggerganov

[spencer-zaid]

fairydreaming#2

Wired the fused CPU implementation and implemented the CUDA kernel for GPU support for personal use. With solid interest from users on reddit, decided to PR it into this branch to bring it in, but let me know if you would prefer I make a separate PR directly into master

https://www.reddit.com/r/LocalLLaMA/comments/1ulymml/llamacpp_patch_deepseek_v4_flash_running_with/

[fairydreaming]

@spencer-zaid Check out #21149, I have my CUDA lightning indexer implementation there.

[ggerganov]

naming (use any abbreviations or not),

Should be OK as it is - already have an op with the same name length.

whether to include scale factors in arguments or leave them outside by prescaling indexer weights,

Should be OK to have the scales as arguments.

whether to perform dot product calculations from fp32 Q and fp32 dequantized K like I did initially or from K and Q quantized to the indexer K type (like in the GGML matrix multiplication implementation),

We can perform the compute in F32 and later extend with ggml_lightning_indexer_set_prec() following the ggml_flash_attn_ext_set_prec() example.

consider adding mask parameter to this OP, so that we don't have to add the mask to the indexer scores (mask could be f16 then).

For consistency with the FA op it's probably worth adding the mask. It can be optional.

We do this for FA and GDN, I'm wondering if we can refactor the code to make it easier to add any op like this

@am17an Yes, we can simplify and streamline the logic for such ops. A solution at the llama.cpp level similar to #24646 should be good.

[fairydreaming]

@ggerganov

whether to include scale factors in arguments or leave them outside by prescaling indexer weights,

Should be OK to have the scales as arguments.

When prescaling indexer weights:

llama.cpp/src/models/deepseek4.cpp

Lines 563 to 565 in 20a04b2

	ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur);
	indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f/sqrtf(float(n_embd_indexer_head*n_indexer_head)));
	cb(indexer_weights, "lid_weights", il);

we only do n_embd_indexer_head * n_ubatch * n_stream multiplications.

If we pass these as indexer arguments then with naive implementation (like the current one) it will be (n_embd_indexer_head + 1) * n_kv * n_ubatch * n_stream multiplications just to scale indexer scores.

That's why I left a note about removing these from arguments, it may be simply more efficient solution - unless there's some specific reason to leave them here. Of course we could prescale indexer weights internally inside the lightning indexer implementation, but that would unnecessarily complicate the implementation.

whether to perform dot product calculations from fp32 Q and fp32 dequantized K like I did initially or from K and Q quantized to the indexer K type (like in the GGML matrix multiplication implementation),

We can perform the compute in F32 and later extend with ggml_lightning_indexer_set_prec() following the ggml_flash_attn_ext_set_prec() example.

It's not about accumulator type, but Q and K types when doing calculations.

In the current CPU implementation I do for each Q K vector pair:

1. K is dequantized to float
2. float K x float Q dot products (both vectors are floats)

If I understand correctly in the CPU flash attention implementation we have:

1. Q is quantized to K type
2. K-typed K x K-typed Q dot products (both vectors have K type)

In the original DeepSeek V3.2 fp8_index_kernel() both Q and K are passed as fp8 and indexer GEMM is performed on fp8 data. That's why I also thought about switching to quantized types for Q x K indexer dot products calculations. In DeepSeek V4 they even switched to fp4 there.

[ggerganov]

If we pass these as indexer arguments then with naive implementation (like the current one) it will be (n_embd_indexer_head + 1) * n_kv * n_ubatch * n_stream multiplications just to scale indexer scores.

I see. So it seems like we want to do the indexer weights scaling before the op and keep the score scale as argument?

In the original DeepSeek V3.2 fp8_index_kernel() both Q and K are passed as fp8 and indexer GEMM is performed on fp8 data. That's why I also thought about switching to quantized types for Q x K indexer dot products calculations. In DeepSeek V4 they even switched to fp4 there.

Sounds like it's safe to use quantized multiplications here. I don't have sense about how computationally expensive this op is. My feeling is keep it simple for now (always convert to F32) and later we can explore quantizing the Q and how much this help the performace.

[fairydreaming]

If we pass these as indexer arguments then with naive implementation (like the current one) it will be (n_embd_indexer_head + 1) * n_kv * n_ubatch * n_stream multiplications just to scale indexer scores.

I see. So it seems like we want to do the indexer weights scaling before the op and keep the score scale as argument?

@ggerganov I think there's some misunderstanding. The current implementation calculates each score like this:

score = 0
for each indexer head {
  qk = ggml_vec_dot_f32(q, k)
  qk *= scale_embd
  score += ReLU(qk) * weight
}
score *= scale_heads

but if you do:

score = 0
for each indexer head {
  qk = ggml_vec_dot_f32(q, k)
  score += ReLU(qk) * weight * scale_embd * scale_heads
}

you get the same result, so by precalculating weight * scale_embd * scale_heads (like we currently do in the code as shown in previous msg) before indexer OP you don't need any scales inside the OP, so I think scale parameters can be removed entirely. Weight values are f32, so I don't think we encounter any numerical precision issues with this?

Edit: Sometimes I wonder what's the point of scaling at all considering the fact that the next operation is ggml_top_k() that doesn't give a damn if you scaled the scores or not.

[ggerganov]

IMO anyway would be fine for now and we can refine later. Ideally, we don't want to break the ggml API by adding/removing arguments, but it seems it is unavoidable for new ops (unless we adopt the process in #24803) because it's hard to foresee all the implications.