CUDA + ggml: add sparse-fa for DSV4/GLM (#27970)

This commit is contained in:
Aman Gupta
2026-09-02 17:27:37 +03:00
committed by GitHub
parent 0f3a71be15
commit 8e93a9773b
12 changed files with 392 additions and 101 deletions
+10 -7
View File
@@ -2540,6 +2540,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
ggml_tensor * kq_mask,
ggml_tensor * sinks,
ggml_tensor * v_mla,
int64_t n_kv_max,
float kq_scale,
int il) const {
const bool v_trans = v->nb[1] > v->nb[2];
@@ -2577,6 +2578,8 @@ ggml_tensor * llm_graph_context::build_attn_mha(
res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});
ggml_flash_attn_ext_add_sinks(cur, sinks);
GGML_ASSERT(n_kv_max >= 0 && n_kv_max <= INT32_MAX);
ggml_flash_attn_ext_set_n_kv_max(cur, static_cast<int32_t>(n_kv_max));
ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);
if (v_mla) {
@@ -2726,7 +2729,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = k_cur;
ggml_tensor * v = v_cur;
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -2825,7 +2828,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (inp->self_v_rot) {
@@ -2916,7 +2919,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -3001,7 +3004,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -3080,7 +3083,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (v_rot) {
@@ -3151,7 +3154,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (k_rot) {
@@ -3210,7 +3213,7 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * k = k_cur;
ggml_tensor * v = v_cur;
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
+1
View File
@@ -1171,6 +1171,7 @@ struct llm_graph_context {
ggml_tensor * kq_mask,
ggml_tensor * sinks, // [n_head_q]
ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
int64_t n_kv_max,
float kq_scale,
int il) const;
+4 -3
View File
@@ -752,7 +752,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention(
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0);
cb(kq_mask, "csa_lid_kq_mask", il);
ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il);
const int64_t n_kv_max = std::min<int64_t>(raw_mask->ne[0], hparams.n_swa) + top_k->ne[0];
ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, n_kv_max, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
@@ -807,7 +808,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention(
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0);
cb(kq_mask, "hca_kq_mask", il);
ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il);
ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
@@ -843,7 +844,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il);
ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
+1 -1
View File
@@ -744,7 +744,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il);
ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il);
cb(cur, "kqv_out", il);
// the rotation is its own inverse, so undo it on the value side of the output