From e84c038310647d1557f2c70951e6d268d624770a Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sun, 26 Jul 2026 16:03:56 +0300 Subject: [PATCH] DS4 optimizations (part 2) (#2179) * DS4 optimizations (part 2) * This is slightly better * Another minor tweak * Increase max. number of graph splitinputs to 64 Else with DS4 we can trun into an assert for specific offload situations with more than one GPU. --- ggml/src/ggml-backend.cpp | 2 +- ggml/src/ggml-cuda/concat.cu | 69 ++++++++++++++- ggml/src/ggml-cuda/dsa_attn.cu | 7 +- ggml/src/ggml.c | 2 +- src/graphs/build_deepseek4.cpp | 150 +++++++++++---------------------- 5 files changed, 120 insertions(+), 110 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 49876eb9..a6a394c4 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1125,7 +1125,7 @@ static bool ggml_is_view_op(enum ggml_op op) { #ifndef GGML_SCHED_MAX_SPLIT_INPUTS // Gemma4 with per-layer embeddings and uses up to 32 inputs -#define GGML_SCHED_MAX_SPLIT_INPUTS 32 +#define GGML_SCHED_MAX_SPLIT_INPUTS 64 #endif #ifndef GGML_SCHED_MAX_COPIES diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu index c8695f2c..9c375d8d 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu @@ -194,6 +194,30 @@ static __global__ void concat_f32_non_cont( } } +static __global__ void k_concat_simple(int64_t n1, int64_t n, const float * __restrict__ src1, const float * __restrict__ src2, + float * __restrict__ dst) { + int64_t i = int64_t(blockIdx.x)*blockDim.x + threadIdx.x; + if (i >= n) { + return; + } + dst[i] = i < n1 ? src1[i] : src2[i - n1]; +} + +static __global__ void k_concat_dim0(int ne0, int ne00, + size_t nb01, size_t nb02, size_t nb03, + size_t nb11, size_t nb12, size_t nb13, + size_t nb1, size_t nb2, size_t nb3, + const float * __restrict__ src1, const float * __restrict__ src2, float * __restrict__ dst) { + + src1 += blockIdx.x * nb01 + blockIdx.y * nb02 + blockIdx.z * nb03; + src2 += blockIdx.x * nb11 + blockIdx.y * nb12 + blockIdx.z * nb13; + dst += blockIdx.x * nb1 + blockIdx.y * nb2 + blockIdx.z * nb3; + + for (int i = threadIdx.x; i < ne0; i += blockDim.x) { + dst[i] = i < ne00 ? src1[i] : src2[i - ne00]; + } +} + void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; @@ -211,13 +235,45 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && (dim == 3 || (dim == 2 && dst->ne[3] == 1) || (dim == 1 && dst->ne[2]*dst->ne[3] == 1))) { - const size_t size0 = ggml_nbytes(src0); - const size_t size1 = ggml_nbytes(src1); - CUDA_CHECK(cudaMemcpyAsync((char *)dst->data, src0->data, size0, cudaMemcpyDeviceToDevice, stream)); - CUDA_CHECK(cudaMemcpyAsync((char *)dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream)); + //printf("%s(%s): using cudaMemcpyAsync\n", __func__, dst->name); + constexpr int k_block_size = 512; + int64_t n1 = ggml_nbytes(src0); + int64_t n2 = ggml_nbytes(src1); + if (n1 % sizeof(float) == 0 && n2 % sizeof(float) == 0) { + n1 /= sizeof(float); + n2 /= sizeof(float); + int64_t n = n1 + n2; + int nblocks = (n + k_block_size - 1)/k_block_size; + k_concat_simple<<>>(n1, n, + (const float *)src0->data, (const float *)src1->data, (float *)dst->data); + return; + } + //const size_t size0 = ggml_nbytes(src0); + //const size_t size1 = ggml_nbytes(src1); + CUDA_CHECK(cudaMemcpyAsync((char *)dst->data, src0->data, n1, cudaMemcpyDeviceToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync((char *)dst->data + n1, src1->data, n2, cudaMemcpyDeviceToDevice, stream)); return; } + if (dim == 0 && src0->nb[0] == ggml_type_size(src0->type) && src1->nb[0] == ggml_type_size(src1->type) && + src0->nb[1] % sizeof(float) == 0 && src1->nb[1] % sizeof(float) == 0) { + auto row_size_src0 = ggml_row_size(dst->type, src0->ne[0]); + auto row_size_src1 = ggml_row_size(dst->type, src1->ne[0]); + auto row_size_dst = ggml_row_size(dst->type, dst->ne[0]); + if (row_size_src0 % sizeof(float) == 0 && row_size_src1 % sizeof(float) == 0 && row_size_dst % sizeof(float) == 0) { + auto ne00_eff = row_size_src0/sizeof(float); + auto ne10_eff = row_size_src1/sizeof(float); + auto ne0_eff = row_size_dst /sizeof(float); + dim3 grid(dst->ne[1], dst->ne[2], dst->ne[3]); + k_concat_dim0<<>>(ne0_eff, ne00_eff, + src0->nb[1]/sizeof(float), src0->nb[2]/sizeof(float), src0->nb[3]/sizeof(float), + src1->nb[1]/sizeof(float), src1->nb[2]/sizeof(float), src1->nb[3]/sizeof(float), + dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), dst->nb[3]/sizeof(float), + (const float *)src0->data, (const float *)src1->data, (float *)dst->data); + return; + } + } + if (dim == 0 && src0->nb[0] == ggml_type_size(src0->type) && src1->nb[0] == ggml_type_size(src1->type) && src0->nb[1] % sizeof(float) == 0 && src1->nb[1] % sizeof(float) == 0) { auto bs = ggml_blck_size(dst->type); @@ -225,6 +281,7 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { auto ne00_eff = (src0->ne[0]/bs)*ts/sizeof(float); auto ne0_eff = (dst->ne[0]/bs)*ts/sizeof(float); if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) { + //printf("%s(%s): using dim0 contiguous float version with ne3 = %ld\n", __func__, dst->name, dst->ne[3]); //if (dst->ne[1] >= 65536 || dst->ne[2] >= 65536) { // fprintf(stderr, "%s: ne1 = %ld, ne2 = %ld exceed max. blocks when computing %s\n", __func__, dst->ne[1], dst->ne[2], dst->name); // GGML_ABORT("fatal error"); @@ -249,6 +306,7 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { //printf("%s(not contiguous): %s(%s) and %s(%s)\n", __func__, src0->name, ggml_type_name(src0->type), src1->name, ggml_type_name(src1->type)); auto ne10_eff = (src1->ne[0]/bs)*ts/sizeof(float); dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); + //printf("%s(%s): using dim0 non-contiguous float version\n", __func__, dst->name); concat_f32_non_cont<<>>( (const char *)src0->data, (const char *)src1->data, @@ -271,6 +329,7 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(dst->type == GGML_TYPE_F32); if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && dim == 2 && dst->ne[3] > 1 && src1->ne[2] == 1) { + //printf("%s(%s): using contiguous dim2 float\n", __func__, dst->name); float * dst_d = (float *)dst->data; float * src0_d = (float *)src0->data; float * src1_d = (float *)src1->data; @@ -279,6 +338,7 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { } if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) { + printf("%s(%s): using generic contiguous dim2 float\n", __func__, dst->name); //if (dst->ne[1] >= 65536 || dst->ne[2] >= 65536) { // fprintf(stderr, "%s: ne1 = %ld, ne2 = %ld exceed max. blocks when computing %s\n", __func__, dst->ne[1], dst->ne[2], dst->name); // GGML_ABORT("fatal error"); @@ -297,6 +357,7 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { dst->ne[0], dst->ne[1], dst->ne[2], dim, stream); } } else { + printf("%s(%s): using generic non-contiguous dim2 float\n", __func__, dst->name); dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); concat_f32_non_cont<<>>( (const char *)src0->data, diff --git a/ggml/src/ggml-cuda/dsa_attn.cu b/ggml/src/ggml-cuda/dsa_attn.cu index 8dc9c04d..17b6b6ee 100644 --- a/ggml/src/ggml-cuda/dsa_attn.cu +++ b/ggml/src/ggml-cuda/dsa_attn.cu @@ -221,7 +221,12 @@ bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) if (indexer->ne[0] % 256 != 0) return false; // lazyness to add checks and handle tails in case of not multiple of 256 // But are there DSA variants where top_k is not a multiple of 256? - if (K->ne[1] < 4*indexer->ne[0]) return false; // for efficiency + //if (K->ne[1] < 4*indexer->ne[0]) return false; // for efficiency + if (Q->ne[1] <= 16) { + if (indexer->ne[0] >= K->ne[1]) return false; + } else { + if (K->ne[1] < 4*indexer->ne[0]) return false; // for efficiency + } if (K->ne[2] > 1 || K->ne[3] > 1 || mask->ne[2] > 1 || mask->ne[3] > 1 || Q->ne[3] > 1) return false; if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16 || mask->type != GGML_TYPE_F16 || Q->type != GGML_TYPE_F32) return false; if (K->ne[0] != Q->ne[0]) return false; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 2d443c94..ee2e3167 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -10680,7 +10680,7 @@ struct ggml_tensor * ggml_flash_attn_ext( // TODO: check if vT can be multiplied by (k*qT) if (mask) { - GGML_ASSERT(ggml_is_contiguous(mask)); + //GGML_ASSERT(ggml_is_contiguous(mask)); GGML_ASSERT(mask->ne[2] == 1); GGML_ASSERT(mask->ne[3] == 1); GGML_ASSERT(mask->ne[1] >= GGML_PAD(q->ne[1], GGML_KQ_MASK_PAD) && diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index 42e36338..697ead0d 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -84,7 +84,6 @@ static void dsv4_build_plan_inputs( const char * tag, int64_t n_tokens, bool create_mask = true, bool flash_attn = true) { - //printf("%s(%s): n_tokens = %ld\n", __func__, tag, n_tokens); dsv4_new_i32_input(ctx, &inputs.state_pos, (int64_t) plan.state_pos.size(), (std::string(tag) + "_state_pos").c_str()); dsv4_new_i32_input(ctx, &inputs.state_persist_src_idxs, (int64_t) plan.state_persist_src_idxs.size(), (std::string(tag) + "_persist_src").c_str()); dsv4_new_i32_input(ctx, &inputs.state_persist_dst_idxs, (int64_t) plan.state_persist_dst_idxs.size(), (std::string(tag) + "_persist_dst").c_str()); @@ -99,11 +98,14 @@ static void dsv4_build_plan_inputs( } } -static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, bool neg_inf) { - ggml_tensor * row = ggml_view_1d(ctx, t, t->ne[0], 0); - row = neg_inf ? ggml_scale_bias(ctx, row, 0.0f, -INFINITY) : ggml_scale(ctx, row, 0.0f); - row = ggml_reshape_2d(ctx, row, t->ne[0], 1); - return dsv4_concat_named(ctx, t, row, 1, "dsv4_append_zero_row"); +static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, ggml_tensor **append_row, bool neg_inf) { + if (*append_row == nullptr) { + ggml_tensor * row = ggml_view_1d(ctx, t, t->ne[0], 0); + row = neg_inf ? ggml_scale_bias(ctx, row, 0.0f, -INFINITY) : ggml_scale(ctx, row, 0.0f); + row = ggml_reshape_2d(ctx, row, t->ne[0], 1); + *append_row = row; + } + return dsv4_concat_named(ctx, t, *append_row, 1, "dsv4_append_zero_row"); } static ggml_tensor * dsv4_cache_view_2d( @@ -155,7 +157,6 @@ static ggml_tensor * dsv4_build_raw_mask_view( } if (n_stream <= 0 || n_tokens % n_stream != 0 || raw_k_read_idxs->ne[0] < n_rows_stream*n_stream) { - //printf("%s(Oops): %d, %d, %d\n", __func__, n_stream <= 0, n_tokens % n_stream != 0, raw_k_read_idxs->ne[0] < n_rows_stream*n_stream); ggml_tensor * base = ggml_cont(ctx, ggml_view_2d(ctx, mask, n_kv, n_tokens, mask->nb[1], 0)); cb(base, "mask_base1", il); return dsv4_build_mask_stream_view(ctx, base, std::max(1, n_stream), n_tokens); @@ -194,8 +195,6 @@ static ggml_tensor * dsv4_pad_raw_k_to( return raw_k; } - //printf("Oops: padding KV cache\n"); - const int64_t n_pad = n_kv_target - n_kv_cur; ggml_tensor * row0 = ggml_view_4d(ctx, raw_k, raw_k->ne[0], raw_k->ne[1], 1, raw_k->ne[3], @@ -486,23 +485,6 @@ static ggml_tensor * dsv4_repeat_streams(ggml_context * ctx, ggml_tensor * t, in 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) { - GGML_UNUSED(ctx); - GGML_UNUSED(n_head); - - if (!cparams.flash_attn || kq_mask->ne[3] == 1) { - return nullptr; - } - - // The zero-bias fallback is only needed for unified multi-stream KV. - // The DSV4 cache/controller is non-unified, so keep the direct FA path. - return nullptr; -} - static ggml_tensor * dsv4_build_attn( ggml_context * ctx, const llama_hparams & hparams, @@ -510,7 +492,6 @@ static ggml_tensor * dsv4_build_attn( ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, - ggml_tensor * kq_b, ggml_tensor * kq_mask, ggml_tensor * sinks, float kq_scale, @@ -522,7 +503,6 @@ static ggml_tensor * dsv4_build_attn( const int64_t n_stream = k->ne[3]; if (!cparams.flash_attn && n_stream > 1) { - GGML_ASSERT(kq_b == nullptr); GGML_ASSERT(q->ne[2] % n_stream == 0); const int64_t n_tokens_stream = q->ne[2]/n_stream; ggml_tensor * result = nullptr; @@ -544,7 +524,7 @@ static ggml_tensor * dsv4_build_attn( } ggml_tensor * cur_s = dsv4_build_attn(ctx, hparams, cparams, - q_s, k_s, v_s, nullptr, mask_s, sinks, kq_scale, cb, il, n_compressed, gf); + q_s, k_s, v_s, mask_s, sinks, kq_scale, cb, il, n_compressed, gf); result = result == nullptr ? cur_s : ggml_concat(ctx, result, cur_s, 1); } return result; @@ -561,10 +541,8 @@ static ggml_tensor * dsv4_build_attn( // through Flash Attention accidentally. constexpr bool kv_unified = false; const bool use_flash_attn = cparams.flash_attn && - (!kv_unified || kq_mask->ne[3] == 1) && - kq_b == nullptr; + (!kv_unified || kq_mask->ne[3] == 1); if (use_flash_attn) { - GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet"); if (v_trans) { v = ggml_transpose(ctx, v); @@ -589,25 +567,6 @@ static ggml_tensor * dsv4_build_attn( selected = ggml_mask_to_index(ctx, kq_mask, n_compressed_padded); cb(selected, "mask_to_idx", il); ggml_build_forward_expand(gf, selected); - //if (q->ne[1] == 1) { - // selected = ggml_view_1d(ctx, selected, selected->ne[0], 0); - // kq_mask = ggml_view_1d(ctx, kq_mask, kq_mask->ne[0], 0); - // kq_mask = ggml_reshape_2d(ctx, kq_mask, 1, kq_mask->ne[0]); - // kq_mask = ggml_get_rows(ctx, kq_mask, selected); - // kq_mask = ggml_reshape_1d(ctx, kq_mask, kq_mask->ne[1]); - // k = ggml_get_rows(ctx, k, selected); - // auto kq = ggml_mul_mat(ctx, k, q); - // if (kq_b != nullptr) { - // kq = ggml_add(ctx, kq, kq_b); - // } - // kq = ggml_soft_max_ext(ctx, kq, kq_mask, kq_scale, 0.0f); - // ggml_soft_max_add_sinks(kq, sinks); - // v = ggml_cont(ctx, ggml_transpose(ctx, k)); - // auto kqv = ggml_mul_mat(ctx, v, kq); - // kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3); - // kqv = ggml_reshape_2d(ctx, kqv, kqv->ne[0]*kqv->ne[1], kqv->ne[2]*kqv->ne[3]); - // return kqv; - //} } } @@ -630,11 +589,6 @@ static ggml_tensor * dsv4_build_attn( cb(kq, "kq", il); ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - if (kq_b != nullptr) { - kq = ggml_add(ctx, kq, kq_b); - cb(kq, "kq_plus_kq_b", il); - } - if (kq->type != GGML_TYPE_F32) { kq = ggml_cast(ctx, kq, GGML_TYPE_F32); } @@ -846,11 +800,6 @@ static ggml_tensor * build_overlap_compressed_kv_from_state( GGML_ASSERT(n_blocks > 0); GGML_ASSERT(state_read_idxs != nullptr); - // TODO: remove this. With a specialized op we can store -1 into the index for negative positions - // and then set the appropriate values (0 or inf) in the kernel. - kv_state = dsv4_append_zero_row(ctx0, kv_state, false); - score_state = dsv4_append_zero_row(ctx0, score_state, true); - ggml_tensor * comp = ggml_ds4_comp(ctx0, kv_state, score_state, state_read_idxs, ratio, 0); llm.cb(comp, tag, il); @@ -1085,6 +1034,11 @@ ggml_cgraph * llm_build_context::build_deepseek4() { inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); cb(inpL, "hc_init", -1); + ggml_tensor * append_csa_state = nullptr; + ggml_tensor * append_csa_score = nullptr; + ggml_tensor * append_lid_state = nullptr; + ggml_tensor * append_lid_score = nullptr; + for (int il = 0; il < n_layer; ++il) { ggml_tensor * residual = inpL; ggml_tensor * post = nullptr; @@ -1192,6 +1146,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() { csa_state_score = ggml_add(ctx0, csa_state_score, csa_ape_rows); ggml_tensor * csa_dep = nullptr; + csa_state_kv = dsv4_append_zero_row(ctx0, csa_state_kv, &append_csa_state, false); + csa_state_score = dsv4_append_zero_row(ctx0, csa_state_score, &append_csa_score, true); + if (lctx.dsv4.inputs.csa.state_write_idxs != nullptr && lctx.dsv4.csa_plan.state_write_idxs.size() > 0) { ggml_tensor * csa_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.csa_state_kv[il], csa_state_kv, 1, "dsv4_csa_source_kv"); ggml_tensor * csa_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.csa_state_score[il], csa_state_score, 1, "dsv4_csa_source_score"); @@ -1235,6 +1192,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() { lid_state_score = ggml_add(ctx0, lid_state_score, lid_ape_rows); ggml_tensor * lid_dep = nullptr; + lid_state_kv = dsv4_append_zero_row(ctx0, lid_state_kv, &append_lid_state, false); + lid_state_score = dsv4_append_zero_row(ctx0, lid_state_score, &append_lid_score, true); + if (lctx.dsv4.inputs.lid.state_write_idxs != nullptr && lctx.dsv4.lid_plan.state_write_idxs.size() > 0) { ggml_tensor * lid_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.lid_state_kv[il], lid_state_kv, 1, "dsv4_lid_source_kv"); ggml_tensor * lid_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.lid_state_score[il], lid_state_score, 1, "dsv4_lid_source_score"); @@ -1352,6 +1312,20 @@ ggml_cgraph * llm_build_context::build_deepseek4() { cb(raw_mask, "dsv4_raw_mask_padded", il); ggml_tensor * attn = nullptr; + if (hparams.n_swa > 0) { + constexpr int k_fa_chunk = 256; + int n_swa = hparams.n_swa; + int ntokens = std::max(k_fa_chunk, int(q->ne[2])); + int nton = k_fa_chunk*((ntokens + n_swa + k_fa_chunk - 1)/k_fa_chunk); + int first = raw_k->ne[2] - nton; + if (first > 0) { + raw_k = ggml_view_4d(ctx0, raw_k, raw_k->ne[0], raw_k->ne[1], nton, raw_k->ne[3], + raw_k->nb[1], raw_k->nb[2], raw_k->nb[3], raw_k->nb[2]*first); + raw_mask = ggml_view_4d(ctx0, raw_mask, nton, raw_mask->ne[1], raw_mask->ne[2], raw_mask->ne[3], + raw_mask->nb[1], raw_mask->nb[2], raw_mask->nb[3], raw_mask->nb[0]*first); + } + } + if (ratio == llama_context::dsv4_runtime::CSA_RATIO && lctx.dsv4.inputs.csa.kq_mask != nullptr && lctx.dsv4.csa_plan.n_kv > 0 && @@ -1362,12 +1336,15 @@ ggml_cgraph * llm_build_context::build_deepseek4() { lctx.dsv4.csa_ctx, n_embd_head, lctx.dsv4.cache.csa_k[il]->ne[1]/std::max(1, lctx.dsv4.cache.n_stream)); - ggml_tensor * top_k = dsv4_build_lid_top_k(ctx0, *this, qr, cur, inp_pos, il, gf, cb); - ggml_tensor * csa_mask = build_top_k_mask(ctx0, - dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.csa.kq_mask, nullptr, + auto csa_mask = lctx.dsv4.inputs.csa.kq_mask; + if (hparams.indexer_top_k < lctx.dsv4.inputs.csa.kq_mask->ne[0]) { + auto top_k = dsv4_build_lid_top_k(ctx0, *this, qr, cur, inp_pos, il, gf, cb); + csa_mask = build_top_k_mask(ctx0, + dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.csa.kq_mask, nullptr, lctx.dsv4.csa_plan.n_kv, n_tokens, csa_k->ne[3], cb, il), - top_k); - cb(csa_mask, "csa_mask", il); + top_k); + cb(csa_mask, "csa_mask", il); + } const bool use_fattn = cparams.flash_attn; if (use_fattn) { csa_mask = dsv4_pad_mask_tokens(ctx0, csa_mask, n_tokens); @@ -1384,35 +1361,16 @@ ggml_cgraph * llm_build_context::build_deepseek4() { if (raw_mask->type != csa_mask->type) { raw_mask = ggml_cast(ctx0, raw_mask, csa_mask->type); } - { - constexpr int k_fa_chunk = 256; - int n_swa = hparams.n_swa; - int ntokens = std::max(k_fa_chunk, int(q->ne[2])); - int nton = k_fa_chunk*((ntokens + n_swa + k_fa_chunk - 1)/k_fa_chunk); - int first = raw_k->ne[2] - nton; - if (first > 0) { - raw_k = ggml_view_4d(ctx0, raw_k, raw_k->ne[0], raw_k->ne[1], nton, raw_k->ne[3], - raw_k->nb[1], raw_k->nb[2], raw_k->nb[3], raw_k->nb[2]*first); - raw_mask = ggml_view_4d(ctx0, raw_mask, nton, raw_mask->ne[1], raw_mask->ne[2], raw_mask->ne[3], - raw_mask->nb[1], raw_mask->nb[2], raw_mask->nb[3], raw_mask->nb[0]*first); - } - } if (raw_k->type != csa_k->type) { csa_k = ggml_cast(ctx0, csa_k, raw_k->type); } ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2); - //printf("k_all: %ld x %ld x %ld x %ld, raw_k: %ld x %ld x %ld x %ld, csa_k = %ld x %ld x %ld x %ld, q = %ld x %ld x %ld x %ld\n", - // k_all->ne[0], k_all->ne[1], k_all->ne[2], k_all->ne[3], - // raw_k->ne[0], raw_k->ne[1], raw_k->ne[2], raw_k->ne[3], - // csa_k->ne[0], csa_k->ne[1], csa_k->ne[2], csa_k->ne[3], - // q->ne[0], q->ne[1], q->ne[2], q->ne[3]); ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); - ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); cb(csa_k, "csa_k", il); cb(k_all, "csa_k_all", il); cb(kq_mask, "csa_kq_mask", il); int n_csa = hparams.n_swa + hparams.indexer_top_k; - attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_b, kq_mask, model.layers[il].attn_sinks, kq_scale, cb, il, n_csa, gf); + attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_mask, model.layers[il].attn_sinks, kq_scale, cb, il, n_csa, gf); cb(attn, "attn_csa", il); } else if (ratio == llama_context::dsv4_runtime::HCA_RATIO && lctx.dsv4.inputs.hca.kq_mask != nullptr && @@ -1436,35 +1394,21 @@ ggml_cgraph * llm_build_context::build_deepseek4() { if (raw_mask->type != hca_mask->type) { raw_mask = ggml_cast(ctx0, raw_mask, hca_mask->type); } - { - constexpr int k_fa_chunk = 256; - int n_swa = hparams.n_swa; - int ntokens = std::max(k_fa_chunk, int(q->ne[2])); - int nton = k_fa_chunk*((ntokens + n_swa + k_fa_chunk - 1)/k_fa_chunk); - int first = raw_k->ne[2] - nton; - if (first > 0) { - raw_k = ggml_view_4d(ctx0, raw_k, raw_k->ne[0], raw_k->ne[1], nton, raw_k->ne[3], - raw_k->nb[1], raw_k->nb[2], raw_k->nb[3], raw_k->nb[2]*first); - raw_mask = ggml_view_4d(ctx0, raw_mask, nton, raw_mask->ne[1], raw_mask->ne[2], raw_mask->ne[3], - raw_mask->nb[1], raw_mask->nb[2], raw_mask->nb[3], raw_mask->nb[0]*first); - } - } if (hca_k->type != raw_k->type) { hca_k = ggml_cast(ctx0, hca_k, raw_k->type); } ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2); ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); - ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); cb(hca_k, "hca_k", il); cb(k_all, "hca_k_all", il); cb(kq_mask, "hca_kq_mask", il); int n_hca = (n_kv + llama_context::dsv4_runtime::HCA_RATIO - 1)/llama_context::dsv4_runtime::HCA_RATIO; n_hca += hparams.n_swa; - attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_b, kq_mask, model.layers[il].attn_sinks, kq_scale, cb, il, n_hca, gf); + attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_mask, model.layers[il].attn_sinks, kq_scale, cb, il, n_hca, gf); cb(attn, "attn_hca", il); } else { - ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, raw_mask, q->ne[1]); - attn = dsv4_build_attn(ctx0, hparams, cparams, q, raw_k, raw_k, kq_b, raw_mask, model.layers[il].attn_sinks, kq_scale, cb, il, -1, gf); + //printf("Regular attention for layer %d\n", il); + attn = dsv4_build_attn(ctx0, hparams, cparams, q, raw_k, raw_k, raw_mask, model.layers[il].attn_sinks, kq_scale, cb, il, -1, gf); cb(attn, "attn_raw", il); }