CUDA + ggml: add sparse-fa for DSV4/GLM (#27970)
This commit is contained in:
+10
-7
@@ -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) {
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user