diff --git a/common/common.cpp b/common/common.cpp index b1fa0f12a..f511aeb14 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/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 dfe4eca75..67e381b57 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -490,7 +490,7 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer( cb(indexer_score, "indexer_score_f32", il); } - if (cparams.flash_attn) { + 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); 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 c99b86874..ea5c5c5a5 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -6641,6 +6641,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, @@ -7060,6 +7061,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, @@ -7211,6 +7213,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);