ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer (#24231)
* ggml : add GGML_OP_LIGHTNING_INDEXER that implements DeepSeek V3.2/V4 lightning indexer * ggml : remove scale parameters from lightning indexer OP, add f16 mask parameter * tests : add GGML_OP_LIGHTNING_INDEXER tests * ggml : bump RPC version * chore : check if lightning indexer input tensors are not transposed * tests : count flops instead of bandwidth in lightning indexer test * chore : add missing const * chore : whitespace * ggml : renamed variables in CPU lightning indexer implementation * ggml : fix lightning indexer mask broadcasting * tests : tests for lightning indexer mask broadcasting * chore : whitespace * llama : use GGML_OP_LIGHTNING_INDEXER in DeepSeek V3.2 and DeepSeek V4 models --------- Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
This commit is contained in:
co-authored by
Stanisław Szymczyk
parent
76f2798059
commit
00f5442cc4
+22
-15
@@ -556,25 +556,32 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(
|
||||
indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream,
|
||||
indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
|
||||
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "lid_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "lid_k", il);
|
||||
ggml_tensor * indexer_score = nullptr;
|
||||
if (cparams.fused_lid) {
|
||||
indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask);
|
||||
cb(indexer_score, "lid_score_masked", il);
|
||||
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
|
||||
} else {
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "lid_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "lid_k", il);
|
||||
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "lid_kq", il);
|
||||
|
||||
ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "lid_score", il);
|
||||
indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "lid_score", il);
|
||||
|
||||
indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
|
||||
cb(indexer_score, "lid_score_masked", il);
|
||||
indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);
|
||||
cb(indexer_score, "lid_score_masked", il);
|
||||
}
|
||||
|
||||
const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k;
|
||||
ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));
|
||||
|
||||
Reference in New Issue
Block a user