nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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.
This commit is contained in:
Kawrakow 2026-07-26 16:03:56 +03:00 committed by GitHub
parent b20bff2ae0
commit e84c038310
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 120 additions and 110 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ggml/src/ggml-backend.cpp
View File

@ -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

69
ggml/src/ggml-cuda/concat.cu
View File

@ -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<<<nblocks, k_block_size, 0, ctx.stream()>>>(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<<<grid, CUDA_CONCAT_BLOCK_SIZE, 0, ctx.stream()>>>(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<<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
(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<<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
(const char *)src0->data,

7
ggml/src/ggml-cuda/dsa_attn.cu
View File

@ -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;

2
ggml/src/ggml.c
View File

@ -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) &&

150
src/graphs/build_deepseek4.cpp
View File

@ -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<int64_t>(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<uint32_t>(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);
}