Make it a command line option

This commit is contained in:
Kawrakow
2026-07-08 06:58:50 +00:00
parent ca03df7f1e
commit bdb927cf4b
6 changed files with 19 additions and 1 deletions
+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;
+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;
+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;
+1 -1
View File
@@ -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);
+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;
+9
View File
@@ -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);