furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 02:05:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fused indexer top_k (CPU only) (#2098)

Browse Source 
* 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
This commit is contained in:
Kawrakow
2026-07-09 18:04:41 +03:00
committed by GitHub
parent c2b58f88dd
commit 6d30fa2fe6
10 changed files with 288 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+6
View File
@@ -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;
common/common.h
+1
View File
@@ -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;
ggml/include/ggml.h
+10
View File
@@ -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 *);
ggml/src/ggml.c
+45 -3
View File
@@ -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);
ggml/src/iqk/iqk_mul_mat.cpp
+202
View File
@@ -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
ggml/src/iqk/iqk_mul_mat.h
+3
View File
@@ -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
include/llama.h
+1
View File
@@ -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;
src/graphs/build_deepseek2.cpp
+10
View File
@@ -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.
src/llama-cparams.h
+1
View File
@@ -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;
src/llama.cpp
+9
View File
@@ -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);