From 2ed3c1abbb8e155226b0b2cbeb9e9efad77fbb02 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Fri, 10 Jul 2026 09:06:58 +0200 Subject: [PATCH] llama : make all KQ masks f16 if FA is used, remove zero attention bias, remove raw_k repeats in DeepSeek V4 (#25370) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * llama : make all KQ masks (except the lightning indexer one) f16 if FA is used and remove zero attention bias in DeepSeek V4 * llama : remove dead code that repeats unified raw_k cache for each stream in DeepSeek V4 - no longer needed as raw_k is always non-unified. --------- Co-authored-by: Stanisław Szymczyk --- src/llama-graph.cpp | 38 ++++++++++++++++++--------- src/models/deepseek4.cpp | 56 +++------------------------------------- 2 files changed, 30 insertions(+), 64 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8c9228b38f..a8fd11ebc3 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -646,7 +646,7 @@ static void dsv4_set_kq_mask( return; } - GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); GGML_ASSERT(dst->ne[0] == plan.n_kv); @@ -656,13 +656,27 @@ static void dsv4_set_kq_mask( GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) n_tokens); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); - float * data = (float *) dst->data; + if (dst->type == GGML_TYPE_F32) { + float * data = (float *) dst->data; - for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { - const int32_t n_visible = plan.n_visible[i]; + for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { + const int32_t n_visible = plan.n_visible[i]; - for (int64_t j = 0; j < dst->ne[0]; ++j) { - data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; + for (int64_t j = 0; j < dst->ne[0]; ++j) { + data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; + } + } + } else if (dst->type == GGML_TYPE_F16) { + ggml_fp16_t * data = (ggml_fp16_t *) dst->data; + const ggml_fp16_t fp16_ninf = llama_cast(-INFINITY); + const ggml_fp16_t fp16_zero = llama_cast(0.0f); + + for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { + const int32_t n_visible = plan.n_visible[i]; + + for (int64_t j = 0; j < dst->ne[0]; ++j) { + data[i*dst->ne[0] + j] = j < n_visible ? fp16_zero : fp16_ninf; + } } } } @@ -679,8 +693,7 @@ static ggml_tensor * dsv4_build_raw_kq_mask( GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); - const bool use_fattn = cparams.flash_attn && (!cparams.kv_unified || n_stream == 1); - const auto type = use_fattn ? GGML_TYPE_F16 : GGML_TYPE_F32; + const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(res); @@ -814,6 +827,7 @@ static void dsv4_build_comp_inputs( llm_graph_input_dsv4::comp_input & inp, const llama_kv_cache_dsv4_context::comp_plan & plan, const char * name, + const llama_cparams & cparams, int64_t n_stream) { inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos"); inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs"); @@ -828,7 +842,7 @@ static void dsv4_build_comp_inputs( GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); - inp.kq_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); + inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(inp.kq_mask); ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); } @@ -3076,9 +3090,9 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { inp_raw->self_k_rot = raw_ctx->build_input_k_rot(ctx0); auto inp = std::make_unique(cparams, std::move(inp_raw), mctx_cur); - 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); + dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", cparams, n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", cparams, n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", cparams, 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); diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 3fb5bff1bb..07aa477e1e 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -184,32 +184,6 @@ static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggm return ggml_add(ctx, t, zero); } -// Raw SWA K is stored once, but compressed K/masks can carry a stream axis. -// Repeat raw K at graph build time before concatenating raw and compressed K. -static ggml_tensor * dsv4_repeat_streams(ggml_context * ctx, ggml_tensor * t, int64_t n_stream) { - if (t->ne[3] == n_stream) { - return t; - } - - GGML_ASSERT(t->ne[3] == 1); - return ggml_repeat_4d(ctx, t, t->ne[0], t->ne[1], t->ne[2], n_stream); -} - -static ggml_tensor * dsv4_build_kq_zero_bias( - ggml_context * ctx, - const llama_cparams & cparams, - ggml_tensor * kq_mask, - int64_t n_head) { - if (!cparams.kv_unified || !cparams.flash_attn || kq_mask->ne[3] == 1) { - return nullptr; - } - - // Keep multi-stream unified DSV4 on the explicit attention path. - ggml_tensor * res = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, - kq_mask->ne[0], kq_mask->ne[1], n_head, kq_mask->ne[3]); - return ggml_fill(ctx, res, 0.0f); -} - static constexpr int64_t DSV4_CSA_RATIO = 4; static constexpr int64_t DSV4_HCA_RATIO = 128; @@ -624,7 +598,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask( ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); zeros = ggml_fill(ctx0, zeros, 0.0f); ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); @@ -681,26 +655,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( csa_k->nb[1], csa_k->nb[2], csa_k->nb[3], 0); cb(csa_k, "csa_comp_k", il); - raw_k = dsv4_repeat_streams(ctx0, raw_k, csa_k->ne[3]); - ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2); cb(k_all, "csa_k_all", il); ggml_tensor * raw_mask = inp_attn->get_kq_mask(); ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il); - const bool use_fattn = cparams.flash_attn && (!cparams.kv_unified || csa_mask->ne[3] == 1); - if (use_fattn && csa_mask->type != GGML_TYPE_F16) { - csa_mask = ggml_cast(ctx0, csa_mask, GGML_TYPE_F16); - } - if (raw_mask->type != csa_mask->type) { - raw_mask = ggml_cast(ctx0, raw_mask, csa_mask->type); - } ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); - 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); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -746,26 +710,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( hca_k->nb[1], hca_k->nb[2], hca_k->nb[3], 0); cb(hca_k, "hca_comp_k", il); - raw_k = dsv4_repeat_streams(ctx0, raw_k, hca_k->ne[3]); - ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2); cb(k_all, "hca_k_all", il); ggml_tensor * raw_mask = inp_attn->get_kq_mask(); ggml_tensor * hca_mask = inp_hca.kq_mask; - const bool use_fattn = cparams.flash_attn && (!cparams.kv_unified || hca_mask->ne[3] == 1); - if (use_fattn && hca_mask->type != GGML_TYPE_F16) { - hca_mask = ggml_cast(ctx0, hca_mask, GGML_TYPE_F16); - } - if (raw_mask->type != hca_mask->type) { - raw_mask = ggml_cast(ctx0, raw_mask, hca_mask->type); - } ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); cb(kq_mask, "hca_kq_mask", il); - 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); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -800,10 +754,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * kq_mask = inp_attn->get_kq_mask(); ggml_tensor * k = mctx_cur->get_k(ctx0, il); - k = dsv4_repeat_streams(ctx0, k, kq_mask->ne[3]); - 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); + ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); }