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
@@ -11568,3 +11568,87 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor *
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ggml_compute_forward_lightning_indexer
|
||||
|
||||
void ggml_compute_forward_lightning_indexer(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
|
||||
const ggml_tensor * q = dst->src[0];
|
||||
const ggml_tensor * k = dst->src[1];
|
||||
const ggml_tensor * w = dst->src[2]; // weights
|
||||
const ggml_tensor * m = dst->src[3]; // mask
|
||||
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( q->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( w->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT( m->type == GGML_TYPE_F16);
|
||||
|
||||
GGML_TENSOR_LOCALS(int64_t, neq, q, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbq, q, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, nek, k, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbk, k, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, new, w, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbw, w, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, nem, m, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nbm, m, nb)
|
||||
GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
|
||||
GGML_TENSOR_LOCALS(size_t, nb, dst, nb)
|
||||
|
||||
GGML_ASSERT( nb0 == ggml_type_size(dst->type));
|
||||
GGML_ASSERT(nbq0 == ggml_type_size( q->type));
|
||||
GGML_ASSERT(nbk0 == ggml_type_size( k->type));
|
||||
GGML_ASSERT(nbw0 == ggml_type_size( w->type));
|
||||
GGML_ASSERT(nbm0 == ggml_type_size( m->type));
|
||||
|
||||
const int n_embd = q->ne[0];
|
||||
const int n_head = q->ne[1];
|
||||
const int n_tokens = q->ne[2];
|
||||
const int n_stream = q->ne[3];
|
||||
const int n_kv = k->ne[2];
|
||||
|
||||
ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float;
|
||||
GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type");
|
||||
|
||||
const int nr = n_kv;
|
||||
const int ith = params->ith;
|
||||
const int nth = params->nth;
|
||||
|
||||
// (temporary) buffer for K converted to float
|
||||
float * k_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32);
|
||||
|
||||
// rows per thread
|
||||
const int dr = (nr + nth - 1)/nth;
|
||||
|
||||
// row range for this thread
|
||||
const int ir0 = dr*ith;
|
||||
const int ir1 = MIN(ir0 + dr, nr);
|
||||
|
||||
for (int s = 0; s < n_stream; ++s) {
|
||||
for (int t = 0; t < n_tokens; ++t) {
|
||||
const float * w_row = (float *) ((char *) w->data + t*nbw1 + s*nbw3);
|
||||
const ggml_fp16_t * m_row = (ggml_fp16_t *) ((char *) m->data + t*nbm1 + (s%nem3)*nbm3);
|
||||
float * dst_row = (float *) ((char *) dst->data + t*nb1 + s*nb3 );
|
||||
for (int ik = ir0; ik < ir1; ++ik) {
|
||||
char * k_row = (char *) k->data + ik*nbk2 + s*nbk3;
|
||||
if (k_to_float) {
|
||||
k_to_float(k_row, k_row_f32, n_embd);
|
||||
} else {
|
||||
k_row_f32 = (float *) k_row;
|
||||
}
|
||||
float score = 0.0f;
|
||||
for (int h = 0; h < n_head; ++h) {
|
||||
// dot product of q and k for head h
|
||||
float qk = 0.0f;
|
||||
const float * q_row = (float *) ((char *) q->data + h*nbq1 + t*nbq2 + s*nbq3);
|
||||
ggml_vec_dot_f32(n_embd, &qk, 0, q_row, 0, k_row_f32, 0, 1);
|
||||
// ReLU and weights (prescaled)
|
||||
score += MAX(qk, 0.0f) * w_row[h];
|
||||
}
|
||||
// apply mask
|
||||
dst_row[ik] = score + GGML_CPU_FP16_TO_FP32(m_row[ik]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user