diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index dc41d66904..d78ddf827b 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -63,26 +63,6 @@ static bool can_reuse_kq_mask( // impl -static ggml_tensor * ggml_mul_mat_aux( - ggml_context * ctx, - ggml_tensor * cur, - ggml_tensor * rot) { - const auto n = rot->ne[0]; - - ggml_tensor * res; - - if (!ggml_is_contiguous(cur)) { - res = ggml_cont_2d (ctx, cur, n, ggml_nelements(cur)/n); - } else { - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); - } - res = ggml_mul_mat (ctx, rot, res); - ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); - res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); - - return res; -} - void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) { if (ubatch->token) { const int64_t n_tokens = ubatch->n_tokens; @@ -881,6 +861,14 @@ void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) { dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream); dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream); + if (inp_csa.k_rot && inp_csa.k_rot->buffer) { + mctx->get_csa()->set_input_k_rot(inp_csa.k_rot); + } + + if (inp_hca.k_rot && inp_hca.k_rot->buffer) { + mctx->get_hca()->set_input_k_rot(inp_hca.k_rot); + } + if (inp_lid.k_rot && inp_lid.k_rot->buffer) { mctx->get_lid()->set_input_k_rot(inp_lid.k_rot); } @@ -2633,12 +2621,12 @@ ggml_tensor * llm_graph_context::build_attn( GGML_ASSERT(v_mla == nullptr); if (inp->self_k_rot) { - q_cur = ggml_mul_mat_aux(ctx0, q_cur, inp->self_k_rot); - k_cur = ggml_mul_mat_aux(ctx0, k_cur, inp->self_k_rot); + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); } if (inp->self_v_rot) { - v_cur = ggml_mul_mat_aux(ctx0, v_cur, inp->self_v_rot); + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); } // these nodes are added to the graph together so that they are not reordered @@ -2669,7 +2657,7 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (inp->self_v_rot) { - cur = ggml_mul_mat_aux(ctx0, cur, inp->self_v_rot); + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); } if (wo) { @@ -2874,14 +2862,14 @@ ggml_tensor * llm_graph_context::build_attn( auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot; if (k_rot) { - q_cur = ggml_mul_mat_aux(ctx0, q_cur, k_rot); + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot); if (k_cur) { - k_cur = ggml_mul_mat_aux(ctx0, k_cur, k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot); } } if (v_rot) { if (v_cur) { - v_cur = ggml_mul_mat_aux(ctx0, v_cur, v_rot); + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, v_rot); } } @@ -2924,7 +2912,7 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (v_rot) { - cur = ggml_mul_mat_aux(ctx0, cur, v_rot); + cur = llama_mul_mat_hadamard(ctx0, cur, v_rot); } if (wo) { @@ -3084,6 +3072,8 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", n_stream); dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", n_stream); dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", n_stream); + inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0); + inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0); inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0); return (llm_graph_input_dsv4 *) res->add_input(std::move(inp)); diff --git a/src/llama-impl.h b/src/llama-impl.h index 7923c3f7ed..2d06752c94 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -54,6 +54,26 @@ static inline dst_t llama_cast(src_t v) { } } +static inline ggml_tensor * llama_mul_mat_hadamard( + ggml_context * ctx, + ggml_tensor * cur, + ggml_tensor * rot) { + const auto n = rot->ne[0]; + + ggml_tensor * res; + + if (!ggml_is_contiguous(cur)) { + res = ggml_cont_2d(ctx, cur, n, ggml_nelements(cur)/n); + } else { + res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); + } + res = ggml_mul_mat(ctx, rot, res); + ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); + res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); + + return res; +} + struct time_meas { time_meas(int64_t & t_acc, bool disable = false); ~time_meas(); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 12bf5c3791..680de5144a 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -57,22 +57,6 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { } } -static ggml_tensor * ggml_mul_mat_aux( - ggml_context * ctx, - ggml_tensor * cur, - ggml_tensor * rot) { - const auto n = rot->ne[0]; - - ggml_tensor * res; - - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); - res = ggml_mul_mat (ctx, rot, res); - ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); - res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); - - return res; -} - // // llama_kv_cache // @@ -1875,14 +1859,14 @@ ggml_tensor * llama_kv_cache::build_rope_shift( tmp = ggml_cast(ctx, cur, GGML_TYPE_F32); // rotate back - tmp = ggml_mul_mat_aux(ctx, tmp, rot); + tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_rope_ext(ctx, tmp, shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow); // rotate fwd - tmp = ggml_mul_mat_aux(ctx, tmp, rot); + tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_cpy(ctx, tmp, cur); } else { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 759654228e..3fb5bff1bb 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -557,7 +557,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( cb(indexer_q_pe, "lid_q_pe", il); indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0); - indexer_q = ggml_mul_mat(ctx0, inp_lid.k_rot, indexer_q); + indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot); cb(indexer_q, "lid_q_rot", il); ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur); @@ -652,10 +652,15 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( int il) const { const auto & inp_csa = inp_dsv4->get_csa(); GGML_ASSERT(inp_csa.kq_mask); - GGML_ASSERT(inp_attn->self_k_rot == nullptr); ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } + ggml_build_forward_expand(gf, q); ggml_build_forward_expand(gf, kv); @@ -696,6 +701,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } cb(out, "attn_csa_lid", il); return out; @@ -711,7 +719,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( int il) const { const auto & inp_hca = inp_dsv4->get_hca(); GGML_ASSERT(inp_hca.kq_mask); - GGML_ASSERT(inp_attn->self_k_rot == nullptr); + + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } ggml_build_forward_expand(gf, q); ggml_build_forward_expand(gf, kv); @@ -753,6 +766,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } cb(out, "attn_hca", il); return out; @@ -770,8 +786,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * k_rot = inp_attn->self_k_rot; if (k_rot) { - q = ggml_mul_mat(ctx0, k_rot, q); - kv = ggml_mul_mat(ctx0, k_rot, kv); + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); } ggml_build_forward_expand(gf, q); @@ -788,6 +804,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); ggml_tensor * out = build_attn_mha(q, k, k, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } cb(out, "attn_raw", il); return out; @@ -917,6 +936,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( "csa_state_compress", il); + if (inp_dsv4->get_csa().k_rot) { + kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot); + cb(kv_comp_csa_state, "csa_state_compress_rot", il); + } + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il)); @@ -965,7 +989,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( il); if (inp_dsv4->get_lid().k_rot) { - kv_comp_lid_state = ggml_mul_mat(ctx0, inp_dsv4->get_lid().k_rot, kv_comp_lid_state); + kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot); cb(kv_comp_lid_state, "lid_state_compress_rot", il); } @@ -1007,6 +1031,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( "hca_state_compress", il); + if (inp_dsv4->get_hca().k_rot) { + kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot); + cb(kv_comp_hca, "hca_state_compress_rot", il); + } + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0, kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il)); hca_state_dep = kv_comp_hca; @@ -1035,13 +1064,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( if (ratio == DSV4_CSA_RATIO && inp_dsv4->get_csa().kq_mask && inp_dsv4->get_lid().kq_mask && - inp_dsv4->get_lid().k_rot && - inp_attn->self_k_rot == nullptr) { + inp_dsv4->get_lid().k_rot) { out = build_csa_lid_attention(model, inp_dsv4, inp_attn, q, kv, qr, cur, inp_pos, layer.attn_sinks, 1.0f/sqrtf(float(n_embd_head)), il); } else if (ratio == DSV4_HCA_RATIO && - inp_dsv4->get_hca().kq_mask && - inp_attn->self_k_rot == nullptr) { + inp_dsv4->get_hca().kq_mask) { out = build_hca_attention(inp_dsv4, inp_attn, q, kv, layer.attn_sinks, 1.0f/sqrtf(float(n_embd_head)), il); } else {