fix: GLM-DSA regression from #2067 (server crash with --spec-type mtp) (#2071)

- fix incorrect reshape in the small-batch (n_tokens <= 8) indexer path
- anchor inp_dsa_sink in the graph with ggml_build_forward_expand
This commit is contained in:
usrlocalben 2026-07-03 01:24:06 -05:00 committed by GitHub
parent 87fc8701ff
commit a5e41bc210
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions

View File

@ -475,7 +475,7 @@ ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer(
cb(indexer_kq, "dsa_indexer_kq", il);
indexer_kq = ggml_relu(ctx0, indexer_kq);
cb(indexer_kq, "dsa_indexer_kq_relu", il);
indexer_kq = ggml_reshape_3d(ctx0, indexer_kq, indexer_kq->ne[0], indexer_q->ne[1], indexer_kq->ne[1]/indexer_q->ne[1]); // [n_kv, n_ihead, n_tokens]
indexer_kq = ggml_reshape_3d(ctx0, indexer_kq, indexer_kq->ne[0], n_ihead, n_tokens); // [n_kv, n_ihead, n_tokens]
indexer_kq = ggml_cont(ctx0, ggml_transpose(ctx0, indexer_kq)); // [n_ihead, n_kv, n_tokens]
indexer_weights = ggml_reshape_3d(ctx0, indexer_weights, indexer_weights->ne[0], 1, indexer_weights->ne[1]); // [n_ihead, 1, n_tokens]
indexer_kq = ggml_mul(ctx0, indexer_kq, indexer_weights);
@ -1208,6 +1208,7 @@ ggml_cgraph * llm_build_context::build_deepseek2() {
lctx.inp_dsa_sink = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens);
cb(lctx.inp_dsa_sink, "dsa_sink", -1);
ggml_set_input(lctx.inp_dsa_sink);
ggml_build_forward_expand(gf, lctx.inp_dsa_sink);
}
auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens);
minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY);