From 6d30fa2fe603c02adf7df333639b280c48c5c70f Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Thu, 9 Jul 2026 18:04:41 +0300 Subject: [PATCH] Fused indexer top_k (CPU only) (#2098) * Indexer topk op - CPU only We do save memory, but it is somehow much slower than what we have on main. * Make it a command line option --- common/common.cpp | 6 + common/common.h | 1 + ggml/include/ggml.h | 10 ++ ggml/src/ggml.c | 48 +++++++- ggml/src/iqk/iqk_mul_mat.cpp | 202 +++++++++++++++++++++++++++++++++ ggml/src/iqk/iqk_mul_mat.h | 3 + include/llama.h | 1 + src/graphs/build_deepseek2.cpp | 10 ++ src/llama-cparams.h | 1 + src/llama.cpp | 9 ++ 10 files changed, 288 insertions(+), 3 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index d2a9a549c..387cf280b 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1900,6 +1900,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa params.dsa = true; return true; } + if (arg == "-fidx" || arg == "--fused-indexer-topk") { + params.fused_idx_topk = true; + return true; + } if (arg == "-dsatk" || arg == "--dsa-top-k") { CHECK_ARG params.dsa_top_k = std::stoi(argv[i]); @@ -3025,6 +3029,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", "-fa, --flash-attn (auto|on|off|0|1)", "set Flash Attention (default: %s)", params.flash_attn ? "on" : "off" }); options.push_back({ "*", "-mla, --mla-use", "enable MLA (default: %d)", params.mla_attn }); options.push_back({ "*", "-dsa, --dsa", "enable GLM DSA sparse attention (GLM-DSA arch only; default: %s)", params.dsa ? "enabled" : "disabled" }); + options.push_back({ "*", "-fidx, --fused-indexer-topk", "enable the fused indexer topk op (DSA only; default: %s)", params.fused_idx_topk ? "enabled" : "disabled" }); options.push_back({ "*", "-dsatk, --dsa-top-k", "DSA top-k override; <0 uses the model's configured indexer_top_k (default: %d)", params.dsa_top_k }); options.push_back({ "*", "-amb, --attention-max-batch", "max batch size for attention computations (default: %d)", params.attn_max_batch}); options.push_back({ "*", "-no-fmoe, --no-fused-moe", "disable fused MoE (default: %s)", params.fused_moe_up_gate ? "enabled" : "disabled" }); @@ -4274,6 +4279,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.rope_cache = params.rope_cache; cparams.graph_reuse = params.graph_reuse; cparams.dsa = params.dsa; + cparams.fused_idx_topk = params.fused_idx_topk; cparams.dsa_top_k = params.dsa_top_k; cparams.k_cache_hadamard = params.k_cache_hadamard; cparams.v_cache_hadamard = params.v_cache_hadamard; diff --git a/common/common.h b/common/common.h index 24a041377..2194aef19 100644 --- a/common/common.h +++ b/common/common.h @@ -418,6 +418,7 @@ struct gpt_params { bool rope_cache = false; // if to use RoPE cache (for supported models) bool graph_reuse = true; // if to reuse compute graphs bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa) + bool fused_idx_topk = false; // enable the fused indexer topk op (off by default; opt-in via -fidx pr --fused-indexer-topk) int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k) int min_experts = -1; float thresh_experts = 0; diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index c8c3f4116..a24692edb 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -704,6 +704,7 @@ extern "C" { GGML_OP_FUSED_NORM, GGML_OP_FUSED_RMS_RMS_ADD, GGML_OP_BLEND, + GGML_OP_INDEXER_TOPK, GGML_OP_COUNT, }; @@ -2565,6 +2566,15 @@ extern "C" { struct ggml_tensor * state, struct ggml_tensor * saved_steps); + GGML_API struct ggml_tensor * ggml_indexer_topk( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * q, + struct ggml_tensor * w, + struct ggml_tensor * mask, + enum ggml_unary_op op, + int n_top_k); + // custom operators typedef void (*ggml_unary_op_f32_t) (const int, float *, const float *); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index af2b2fc1c..eb73b4ce6 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4321,9 +4321,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "FUSED_NORM", "FUSED_RMS_RMS_ADD", "BLEND", + "INDEXER_TOPK", }; -static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -4442,10 +4443,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "norm(x,y)", "rms(x1)+rms(x2)", "blend(a,b,c)", + "indexer_topk(k, q, w, mask)", }; -static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -10098,6 +10100,34 @@ struct ggml_tensor * ggml_delta_net( return result; } +struct ggml_tensor * ggml_indexer_topk( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * q, + struct ggml_tensor * w, + struct ggml_tensor * mask, + enum ggml_unary_op op, + int n_top_k) { + GGML_ASSERT(k->ne[2] == 1 && k->ne[3] == 1); + GGML_ASSERT(k->ne[1] > n_top_k); + GGML_ASSERT(k->ne[1] == mask->ne[0]); + GGML_ASSERT(k->ne[0] == q->ne[0]); + GGML_ASSERT(q->ne[2] == mask->ne[1]); + GGML_ASSERT(q->ne[1] == w->ne[0]); + GGML_ASSERT(q->ne[2] == w->ne[1]); + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_top_k, q->ne[2]); + result->op = GGML_OP_INDEXER_TOPK; + result->op_params[0] = (int)op; + result->src[0] = k; + result->src[1] = q; + result->src[2] = w; + result->src[3] = mask; + + return result; +} + + // ggml_fill static struct ggml_tensor * ggml_fill_impl( @@ -24713,6 +24743,12 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml { ggml_compute_forward_delta_net(params, tensor); } break; + case GGML_OP_INDEXER_TOPK: + { + if (!iqk_indexer_topk(tensor, params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, params->ith, params->nth)) { + GGML_ABORT("Fatal error"); + } + } break; case GGML_OP_WIN_PART: { ggml_compute_forward_win_part(params, tensor); @@ -25774,6 +25810,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor case GGML_OP_FILL: case GGML_OP_SOLVE_TRI: case GGML_OP_DELTA_NET: + case GGML_OP_INDEXER_TOPK: { GGML_ABORT("fatal error"); // TODO: not implemented } @@ -26509,6 +26546,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_OUT_PROD: case GGML_OP_SOLVE_TRI: case GGML_OP_DELTA_NET: + case GGML_OP_INDEXER_TOPK: { n_tasks = n_threads; } break; @@ -26833,7 +26871,11 @@ struct ggml_cplan ggml_graph_plan(const struct ggml_cgraph * cgraph, int n_threa cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 } } break; - + case GGML_OP_INDEXER_TOPK: + { + size_t size = iqk_idx_topk_work_buffer_size(node, n_tasks); + cur = MAX(cur, size); + } break; case GGML_OP_CROSS_ENTROPY_LOSS: { cur = ggml_type_size(node->type)*(n_tasks + node->src[0]->ne[0]*n_tasks); diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index ff1149ff8..5158c6c5a 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -1736,6 +1736,201 @@ bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_ty return true; } +namespace { +size_t iqk_idx_topk_work_wbs_per_thread(const struct ggml_tensor * dst, int nth) { + auto k = dst->src[0]; + auto q = dst->src[1]; + if (q->ne[2] >= nth) { + size_t size = 0; + auto tt = ggml_internal_get_type_traits(k->type); + if (tt.is_quantized && tt.vec_dot_type != q->type) { + auto row_size_q = ggml_row_size(tt.vec_dot_type, q->ne[0]); + size = row_size_q * q->ne[1]; + } + size += k->ne[1] * q->ne[1] * sizeof(float); + size += k->ne[1] * sizeof(float); + size += k->ne[1] * sizeof(int32_t); + size = GGML_PAD(size, 128); + return size; + } + return 0; +} +// TODO: SIMDify +inline void iqk_f16_to_f32(int n, const ggml_fp16_t * x, float * y) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_FP16_TO_FP32(x[i]); + } +} +} + +size_t iqk_idx_topk_work_buffer_size(const struct ggml_tensor * dst, int nthread) { + auto k = dst->src[0]; + auto q = dst->src[1]; + if (q->ne[2] >= nthread) { + return iqk_idx_topk_work_wbs_per_thread(dst, nthread) * nthread; + } + size_t size = 0; + auto tt = ggml_internal_get_type_traits(k->type); + if (tt.is_quantized && tt.vec_dot_type != q->type) { + auto row_size_q = ggml_row_size(tt.vec_dot_type, q->ne[0]); + size = row_size_q * q->ne[1] * q->ne[2]; + } + size += k->ne[1] * q->ne[1] * sizeof(float); + size += k->ne[1] * sizeof(float); + size += k->ne[1] * sizeof(int32_t); + return size; +} + +bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t barrier, void * barrier_data, int ith, int nth) { + if (dst->op != GGML_OP_INDEXER_TOPK) { + return false; + } + auto op = ggml_unary_op(dst->op_params[0]); + if (op != GGML_UNARY_OP_RELU) { + return false; + } + int n_top_k = dst->ne[0]; + auto k = dst->src[0]; + auto q = dst->src[1]; + auto w = dst->src[2]; + auto m = dst->src[3]; + if (k->ne[2] != 1 || k->ne[3] != 1) return false; + if (k->ne[1] <= n_top_k ) return false; + if (k->ne[1] != m->ne[0]) return false; + if (k->ne[0] != q->ne[0]) return false; + if (q->ne[2] != m->ne[1]) return false; + if (q->ne[1] != w->ne[0]) return false; + if (q->ne[2] != w->ne[1]) return false; + if (q->type != GGML_TYPE_F32) return false; + if (m->type != GGML_TYPE_F32 && m->type != GGML_TYPE_F16) return false; + + auto work_size = iqk_idx_topk_work_wbs_per_thread(dst, nth); + auto work = (char *)work_buffer + ith*work_size; + ggml_from_float_t from_float = nullptr; + + auto tt = ggml_internal_get_type_traits(k->type); + + size_t quantize_size = 0; + auto q_type = q->type; + auto row_size_q = q->nb[1]; + if (tt.is_quantized && tt.vec_dot_type != q->type) { + auto ttq = ggml_internal_get_type_traits(tt.vec_dot_type); + from_float = ttq.from_float; + row_size_q = ggml_row_size(tt.vec_dot_type, q->ne[0]); + quantize_size = row_size_q * q->ne[1]; + q_type = tt.vec_dot_type; + } + + MulMat mm; + if (!MulMat::prepare(int(k->type), int(q_type), k->ne[0], mm, q->ne[1])) { + return false; + } + + if (q->ne[2] >= nth) { + auto kq = (float *)(work + quantize_size); + auto score = kq + k->ne[1]*q->ne[1]; + auto sorted = (int32_t *)(score + k->ne[1]); + for (int iq = ith; iq < q->ne[2]; iq += nth) { + auto this_q = (const char *)q->data + iq*q->nb[2]; + auto this_m = (const char *)m->data + iq*m->nb[1]; + auto this_w = (const float *)((const char *)w->data + w->nb[1]*iq); + if (from_float) { + from_float((const float *)this_q, work, q->ne[0] * q->ne[1]); + this_q = work; + } + DataInfo info{kq, this_q, (size_t)k->ne[1], (size_t)row_size_q, 0, 1, nullptr, 0}; + mm.mul_mat_NxM(k->ne[0], (const char *)k->data, k->nb[1], info, k->ne[1], q->ne[1]); + auto kq_i = kq; + + if (m->type == GGML_TYPE_F32) { + std::memcpy(score, this_m, k->ne[1]*sizeof(float)); + } else { + iqk_f16_to_f32(k->ne[1], (const ggml_fp16_t *)this_m, score); + } + for (int i = 0; i < int(q->ne[1]); ++i) { + float wi = this_w[i]; + for (int j = 0; j < int(k->ne[1]); ++j) { + float relu = kq_i[j] > 0.0f ? kq_i[j] : 0.0f; + score[j] += wi * relu; + } + kq_i += k->ne[1]; + } + for (int j = 0; j < int(k->ne[1]); ++j) sorted[j] = j; + std::partial_sort(sorted, sorted + n_top_k, sorted + k->ne[1], [score] (int32_t l, int32_t r) -> bool { return score[l] > score[r]; }); + std::memcpy((char *)dst->data + dst->nb[1]*iq, sorted, n_top_k*sizeof(int32_t)); + } + return true; + } + + if (k->ne[1] % 32 != 0) return false; // we can assume cache size is a multiple of at least 32 + + // We have more than 1 thread per q row + // To not make it too complicated, let's just have all threads process the same row + int n32 = k->ne[1] / 32; + int npt = (n32 + nth - 1)/nth; + int ith_mid = nth; + int n_this_thread = npt; + int first = ith*npt; + if (npt*nth > n32) { + ith_mid = n32 - nth*(npt - 1); + if (ith >= ith_mid) { + --n_this_thread; + first = ith_mid*npt + (ith - ith_mid)*n_this_thread; + } + } + n_this_thread *= 32; + first *= 32; + auto q_data = (const char *)q->data; + auto qnb2 = q->nb[2]; + if (from_float) { + auto quantized = (char *)work_buffer; + if (ith < q->ne[2]) { + auto this_q = q_data + ith*q->nb[2]; + from_float((const float *)this_q, quantized + quantize_size*ith, q->ne[0] * q->ne[1]); + } + qnb2 = quantize_size; + q_data = quantized; + work += quantize_size; + + barrier(barrier_data); + } + auto kq = (float *)((char *)work_buffer + quantize_size*q->ne[2]); + auto kq_th = kq + first*q->ne[1]; + auto score = kq + k->ne[1]*q->ne[1]; + auto score_th = score + first; + auto sorted = (int32_t *)(score + k->ne[1]); + for (int iq = 0; iq < q->ne[2]; ++iq) { + auto this_q = q_data + iq*qnb2; + auto this_m = (const char *)m->data + iq*m->nb[1]; + auto this_w = (const float *)((const char *)w->data + w->nb[1]*iq); + DataInfo info{kq_th, this_q, (size_t)n_this_thread, (size_t)row_size_q, 0, 1, nullptr, 0}; + mm.mul_mat_NxM(k->ne[0], (const char *)k->data + first*k->nb[1], k->nb[1], info, n_this_thread, q->ne[1]); + if (m->type == GGML_TYPE_F32) { + std::memcpy(score_th, this_m + first*sizeof(float), n_this_thread*sizeof(float)); + } else { + iqk_f16_to_f32(n_this_thread, (const ggml_fp16_t *)this_m + first, score_th); + } + auto kq_i = kq_th; + for (int i = 0; i < int(q->ne[1]); ++i) { + float wi = this_w[i]; + for (int j = 0; j < n_this_thread; ++j) { + float relu = kq_i[j] > 0.0f ? kq_i[j] : 0.0f; + score_th[j] += wi * relu; + } + kq_i += n_this_thread; + } + barrier(barrier_data); + if (ith == 0) { + for (int j = 0; j < int(k->ne[1]); ++j) sorted[j] = j; + std::partial_sort(sorted, sorted + n_top_k, sorted + k->ne[1], [score] (int32_t l, int32_t r) -> bool { return score[l] > score[r]; }); + std::memcpy((char *)dst->data + dst->nb[1]*iq, sorted, n_top_k*sizeof(int32_t)); + } + barrier(barrier_data); + } + + return true; +} + #else // IQK_IMPLEMENT #include "ggml-impl.h" @@ -1776,5 +1971,12 @@ bool iqk_fused_delta_net(int, int, int, int, int, int, return false; } +bool iqk_indexer_topk(struct ggml_tensor *, void *, barrier_t, void *, int, int) { + return false; +} + +size_t iqk_idx_topk_work_buffer_size(const struct ggml_tensor *, int) { + return 0; +} #endif diff --git a/ggml/src/iqk/iqk_mul_mat.h b/ggml/src/iqk/iqk_mul_mat.h index a51597090..aae2eb197 100644 --- a/ggml/src/iqk/iqk_mul_mat.h +++ b/ggml/src/iqk/iqk_mul_mat.h @@ -78,6 +78,9 @@ IQK_API bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int r const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data, const float * state_in, float * out_data, float * state_out, float * saved_steps, int state_step_stride, int ith, int nth); +IQK_API bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t barrier, void * barrier_data, int ith, int nth); +IQK_API size_t iqk_idx_topk_work_buffer_size(const struct ggml_tensor * dst, int nthread); + #ifdef __cplusplus } #endif diff --git a/include/llama.h b/include/llama.h index fc4e5d470..39a13ee4f 100644 --- a/include/llama.h +++ b/include/llama.h @@ -492,6 +492,7 @@ extern "C" { bool rope_cache; // whether to use RoPE cache [EXPERIMENTAL] bool graph_reuse; // whether to reuse graphs when possible [EXPERIMENTAL] bool dsa; // enable GLM DSA sparse attention (off by default) [EXPERIMENTAL] + bool fused_idx_topk; // enable the fused indexer topk op (off by default) [EXPERIMENTAL] int dsa_top_k; // DSA top-k override (<0 => model's configured indexer_top_k) [EXPERIMENTAL] int min_experts; float thresh_experts; diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index aeb1e6eac..4dab202ee 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -489,6 +489,16 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer( indexer_score = ggml_cast(ctx0, indexer_score, GGML_TYPE_F32); cb(indexer_score, "indexer_score_f32", il); } + + if (cparams.flash_attn && cparams.fused_idx_topk) { + if (lctx.inp_dsa_sink) { + indexer_score = ggml_add(ctx0, indexer_score, lctx.inp_dsa_sink); + cb(indexer_score, "dsa_indexer_score_sink", il); + ggml_build_forward_expand(gf, indexer_score); + } + return ggml_indexer_topk(ctx0, indexer_k_b, indexer_q, indexer_weights, indexer_score, GGML_UNARY_OP_RELU, n_top_k); + } + if (indexer_q->ne[2] <= 8) { // This covers TG and small batches (as needed for instance for speculative decoding). It is quite a bit faster than // the loop over attention heads in the other branch. We limit it to a maximum of 8 tokens to limit compute buffer size. diff --git a/src/llama-cparams.h b/src/llama-cparams.h index ac24ba824..7991df05f 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -43,6 +43,7 @@ struct llama_cparams { bool v_cache_hadamard; bool dsa_indexer_hadamard = true; // apply Walsh-Hadamard rotation to DSA indexer q/k (precision) bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa) + bool fused_idx_topk = false; // enable the fused indexer topk op (off by default; opt-in via -fidx or --fused-indexer-topk) int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k) bool split_mode_graph_scheduling; //bool split_mode_f16; diff --git a/src/llama.cpp b/src/llama.cpp index b49ec6133..e05e95ebf 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -6619,6 +6619,7 @@ struct llama_context_params llama_context_default_params() { /*.rope_cache =*/ false, /*.graph_reuse =*/ true, /*.dsa =*/ false, + /*.fused_idx_topk =*/ false, /*.dsa_top_k =*/ -1, /*.min_experts =*/ -1, /*.thtesh_experts =*/ 0.0f, @@ -7038,6 +7039,7 @@ struct llama_context * llama_init_from_model( cparams.rope_cache = params.rope_cache; cparams.graph_reuse = params.graph_reuse; cparams.dsa = params.dsa; + cparams.fused_idx_topk = params.fused_idx_topk; cparams.dsa_top_k = params.dsa_top_k; // The DSA lightning indexer is built only in the layer-mode (non-TP) attention path. Under // -sm graph / -sm attn the model runs the tensor-parallel attention path, which has no indexer, @@ -7189,6 +7191,13 @@ struct llama_context * llama_init_from_model( LLAMA_LOG_INFO("%s: fused_mmad = %d\n", __func__, cparams.fused_mmad); LLAMA_LOG_INFO("%s: rope_cache = %d\n", __func__, cparams.rope_cache); LLAMA_LOG_INFO("%s: graph_reuse = %d\n", __func__, cparams.graph_reuse); + if (model->arch == LLM_ARCH_GLM_DSA) { + LLAMA_LOG_INFO("%s: dsa = %d\n", __func__, cparams.dsa); + LLAMA_LOG_INFO("%s: dsa_idx_topk = %d\n", __func__, cparams.fused_idx_topk); + if (cparams.dsa_top_k > 0) { + LLAMA_LOG_INFO("%s: dsa_top_k = %d\n", __func__, cparams.dsa_top_k); + } + } LLAMA_LOG_INFO("%s: k_cache_hadam = %d\n", __func__, cparams.k_cache_hadamard); LLAMA_LOG_INFO("%s: v_cache_hadam = %d\n", __func__, cparams.v_cache_hadamard); LLAMA_LOG_INFO("%s: split_mode_graph_scheduling = %d\n", __func__, cparams.split_mode_graph_scheduling);