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

DS4 optimizations (#2169) 

* Adding ds4_comp op with CPU implementation

* ds4_comp on CUDA

* ds4_comp: ratio = 4 specialization

Surprisingly small performance gain

* Also handle HCA via ds4_comp

But much smaller gain, if any.

* Delete commented out stuff

* Remove the [(size_t) il] noise

* Minor

* Fix quantized cache
This commit is contained in:
Kawrakow 2026-07-25 08:52:38 +03:00 committed by GitHub
parent f359df4bc9
commit bd342d624f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 473 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
ggml/include/ggml.h
View File

@ -711,6 +711,7 @@ extern "C" {
GGML_OP_HC_POST,
GGML_OP_MASK_TO_IDX,
GGML_OP_LATENT_ATTN,
GGML_OP_DS4_COMP,
GGML_OP_COUNT,
};
@ -2690,6 +2691,14 @@ extern "C" {
struct ggml_tensor * mask,
int max_row_size);
GGML_API struct ggml_tensor * ggml_ds4_comp(
struct ggml_context * ctx,
struct ggml_tensor * state,
struct ggml_tensor * score,
struct ggml_tensor * idx,
int ratio,
int type);
// custom operators

5
ggml/src/ggml-cuda.cu
View File

@ -60,6 +60,7 @@
#include "ggml-cuda/latent_attn.cuh"
#include "ggml-cuda/blend.cuh"
#include "ggml-cuda/indexer_topk.cuh"
#include "ggml-cuda/ds4_comp.cuh"
#include <algorithm>
#include <array>
@ -4156,6 +4157,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_OP_MASK_TO_IDX:
ggml_cuda_op_mask_to_index(ctx, dst);
break;
case GGML_OP_DS4_COMP:
ggml_cuda_op_ds4_comp(ctx, dst);
break;
default:
return false;
}
@ -5065,6 +5069,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
case GGML_OP_INDEXER_TOPK:
case GGML_OP_MASK_TOPK:
case GGML_OP_MASK_TO_IDX:
case GGML_OP_DS4_COMP:
return true;
case GGML_OP_HC_PRE:
case GGML_OP_HC_POST:

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

@ -199,6 +199,10 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
if (src0->type != src1->type) {
printf("%s: %s is type %s, %s is type %s\n", __func__, src0->name, ggml_type_name(src0->type),
src1->name, ggml_type_name(src1->type));
}
GGML_ASSERT(src0->type == src1->type && src0->type == dst->type);
cudaStream_t stream = ctx.stream();

164
ggml/src/ggml-cuda/ds4_comp.cu Normal file
View File

@ -0,0 +1,164 @@
#include "ds4_comp.cuh"
static __global__ void k_ds4_comp(int ne0, int nblock, int ratio, int nidx,
size_t state_stride, size_t score_stride,
const float * __restrict__ state, const float * __restrict__ score, const int * __restrict__ idx, float * dst) {
int ii = blockIdx.x * blockDim.x + threadIdx.x;
int ib = ii / ne0;
if (ib >= nblock) {
return;
}
int i0 = ii % ne0;
idx += ratio*ib;
int row_p = idx[0];
int row_c = idx[nidx];
float vp = score[row_p*score_stride + i0];
float vc = score[row_c*score_stride + i0 + ne0];
float max_v = max(vp, vc);
for (int ir = 1; ir < ratio; ++ir) {
row_p = idx[ir];
row_c = idx[ir+nidx];
vp = score[row_p*score_stride + i0];
vc = score[row_c*score_stride + i0 + ne0];
max_v = max(max_v, max(vp, vc));
}
float sum_num = 0.0f, sum_den = 0.0f;
for (int ir = 0; ir < ratio; ++ir) {
row_p = idx[ir];
row_c = idx[ir+nidx];
vp = score[row_p*score_stride + i0];
vc = score[row_c*score_stride + i0 + ne0];
float sp = state[row_p*state_stride + i0];
float sc = state[row_c*state_stride + i0 + ne0];
float wp = expf(vp - max_v);
float wc = expf(vc - max_v);
sum_den += wp + wc;
sum_num += wp*sp + wc*sc;
}
dst[ib*ne0 + i0] = sum_num / sum_den;
}
static __global__ void k_ds4_comp_4(int ne0, int nblock, int nidx,
size_t state_stride, size_t score_stride,
const float * __restrict__ state, const float * __restrict__ score, const int * __restrict__ idx, float * dst) {
int ii = blockIdx.x * blockDim.x + threadIdx.x;
int ib = ii / ne0;
if (ib >= nblock) {
return;
}
int i0 = ii % ne0;
idx += 4*ib;
float values[16];
#pragma unroll
for (int ir = 0; ir < 4; ++ir) {
int row_p = idx[ir];
int row_c = idx[ir+nidx];
values[2*ir+0] = score[row_p*score_stride + i0];
values[2*ir+1] = score[row_c*score_stride + i0 + ne0];
values[2*ir+8] = state[row_p*state_stride + i0];
values[2*ir+9] = state[row_c*state_stride + i0 + ne0];
}
float max_v = values[0];
#pragma unroll
for (int ir = 1; ir < 8; ++ir) max_v = max(max_v, values[ir]);
float sum_num = 0.0f, sum_den = 0.0f;
#pragma unroll
for (int ir = 0; ir < 8; ++ir) {
float w = expf(values[ir] - max_v);
sum_den += w;
sum_num += w*values[ir+8];
}
dst[ib*ne0 + i0] = sum_num / sum_den;
}
static __global__ void k_ds4_comp_type1(int ne0, int nblock, int ratio,
size_t state_stride, size_t score_stride,
const float * __restrict__ state, const float * __restrict__ score, const int * __restrict__ idx, float * dst) {
int ii = blockIdx.x * blockDim.x + threadIdx.x;
int ib = ii / ne0;
if (ib >= nblock) {
return;
}
int i0 = ii % ne0;
idx += ratio*ib;
int row = idx[0];
float max_v = score[row*score_stride + i0];
for (int ir = 1; ir < ratio; ++ir) {
row = idx[ir];
float v = score[row*score_stride + i0];
max_v = max(max_v, v);
}
float sum_num = 0.0f, sum_den = 0.0f;
for (int ir = 0; ir < ratio; ++ir) {
row = idx[ir];
float v = score[row*score_stride + i0];
float s = state[row*state_stride + i0];
float w = expf(v - max_v);
sum_den += w;
sum_num += w*s;
}
dst[ib*ne0 + i0] = sum_num / sum_den;
}
void ggml_cuda_op_ds4_comp(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
constexpr int k_block_size = 128;
auto * state = dst->src[0];
auto * score = dst->src[1];
auto * idx = dst->src[2];
GGML_ASSERT(state->type == GGML_TYPE_F32);
GGML_ASSERT(score->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_are_same_shape(score, state));
GGML_ASSERT(state->ne[2] == 1 && state->ne[3] == 1);
GGML_ASSERT( idx->type == GGML_TYPE_I32);
GGML_ASSERT(ggml_nrows(idx) == 1);
int type = dst->op_params[0];
int nblock = dst->ne[1];
if (type == 0) {
GGML_ASSERT(state->ne[0] % 64 == 0);
GGML_ASSERT(dst->ne[0] == state->ne[0]/2);
int ratio = idx->ne[0] / (2*nblock);
GGML_ASSERT(idx->ne[0] % (2*ratio) == 0);
int ne0 = dst->ne[0];
int nelem = ne0 * nblock;
int nb = (nelem + k_block_size - 1)/k_block_size;
if (ratio == 4) {
k_ds4_comp_4<<<nb, k_block_size, 0, ctx.stream()>>>(ne0, nblock, idx->ne[0]/2,
state->nb[1]/sizeof(float), score->nb[1]/sizeof(float),
(const float *)state->data, (const float *)score->data, (const int *)idx->data, (float *)dst->data);
} else {
k_ds4_comp<<<nb, k_block_size, 0, ctx.stream()>>>(ne0, nblock, ratio, idx->ne[0]/2,
state->nb[1]/sizeof(float), score->nb[1]/sizeof(float),
(const float *)state->data, (const float *)score->data, (const int *)idx->data, (float *)dst->data);
}
} else {
GGML_ASSERT(state->ne[0] % 32 == 0);
GGML_ASSERT(dst->ne[0] == state->ne[0]);
int ratio = idx->ne[0] / nblock;
GGML_ASSERT(idx->ne[0] % ratio == 0);
int ne0 = dst->ne[0];
int nelem = ne0 * nblock;
int nb = (nelem + k_block_size - 1)/k_block_size;
k_ds4_comp_type1<<<nb, k_block_size, 0, ctx.stream()>>>(ne0, nblock, ratio,
state->nb[1]/sizeof(float), score->nb[1]/sizeof(float),
(const float *)state->data, (const float *)score->data, (const int *)idx->data, (float *)dst->data);
}
}

3
ggml/src/ggml-cuda/ds4_comp.cuh Normal file
View File

@ -0,0 +1,3 @@
#include "common.cuh"
void ggml_cuda_op_ds4_comp(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

252
ggml/src/ggml.c
View File

@ -4341,9 +4341,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"HC_POST",
"MASK_TO_IDX",
"LATENT_ATTN",
"DS4_COMP",
};
static_assert(GGML_OP_COUNT == 110, "GGML_OP_COUNT != 110");
static_assert(GGML_OP_COUNT == 111, "GGML_OP_COUNT != 111");
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
@ -4469,10 +4470,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"hc_post(x,p,r,c)",
"mask_to_idx(masl)",
"latent_attn_prefix(q,c,pk,pv,mask)",
"ds4_comp(state, score, idx)",
};
static_assert(GGML_OP_COUNT == 110, "GGML_OP_COUNT != 110");
static_assert(GGML_OP_COUNT == 111, "GGML_OP_COUNT != 111");
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
@ -10275,6 +10277,43 @@ struct ggml_tensor * ggml_mask_to_index(
return result;
}
struct ggml_tensor * ggml_ds4_comp(
struct ggml_context * ctx,
struct ggml_tensor * state,
struct ggml_tensor * score,
struct ggml_tensor * idx,
int ratio,
int type) {
GGML_ASSERT(state->type == GGML_TYPE_F32);
GGML_ASSERT(score->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_are_same_shape(score, state));
GGML_ASSERT(state->ne[2] == 1 && state->ne[3] == 1);
GGML_ASSERT( idx->type == GGML_TYPE_I32);
GGML_ASSERT(ggml_nrows(idx) == 1);
int ne0, nblock;
if (type == 0) {
GGML_ASSERT(idx->ne[0] % (2*ratio) == 0);
GGML_ASSERT(state->ne[0] % 64 == 0);
nblock = idx->ne[0] / (2*ratio);
ne0 = state->ne[0]/2;
} else {
GGML_ASSERT(idx->ne[0] % ratio == 0);
GGML_ASSERT(state->ne[0] % 32 == 0);
nblock = idx->ne[0] / ratio;
ne0 = state->ne[0];
}
struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, nblock);
result->src[0] = state;
result->src[1] = score;
result->src[2] = idx;
result->op = GGML_OP_DS4_COMP;
result->op_params[0] = type;
return result;
}
// ggml_fill
@ -24056,6 +24095,205 @@ static void ggml_compute_forward_mask_to_idx(const struct ggml_compute_params *
}
}
static void ggml_compute_forward_ds4_comp_type0(const struct ggml_compute_params * params,
struct ggml_tensor * dst) {
struct ggml_tensor * state = dst->src[0];
struct ggml_tensor * score = dst->src[1];
struct ggml_tensor * idx = dst->src[2];
GGML_ASSERT(state->type == GGML_TYPE_F32);
GGML_ASSERT(score->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_are_same_shape(score, state));
GGML_ASSERT(state->ne[2] == 1 && state->ne[3] == 1);
GGML_ASSERT(state->ne[0] % 64 == 0);
GGML_ASSERT(dst->ne[0] == state->ne[0]/2);
GGML_ASSERT( idx->type == GGML_TYPE_I32);
GGML_ASSERT(ggml_nrows(idx) == 1);
int nblock = dst->ne[1];
int ratio = idx->ne[0] / (2*nblock);
GGML_ASSERT(idx->ne[0] % (2*ratio) == 0);
int ith = params->ith;
int nth = params->nth;
int ne0 = dst->ne[0];
int n32 = ne0/32;
int nchunk = n32*nblock;
int npt = (nchunk + nth - 1)/nth;
int first = ith*npt;
int last = MIN(first + npt, nchunk);
const int * index = (const int *)idx->data;
size_t work_size = 32*(4*ratio + 3)*sizeof(float);
GGML_ASSERT(nth*work_size <= params->wsize);
float * work = (float *)((char *)params->wdata + ith*work_size);
float * max_l = work;
float * sum_l = max_l + 32;
float * res_l = sum_l + 32;
float * score_l = res_l + 32;
float * state_l = score_l + 64*ratio;
for (int ic = first; ic < last; ++ic) {
int ib = ic / n32;
int i32 = ic - ib*n32;
int first_i0 = 32*i32;
for (int ir = 0; ir < ratio; ++ir) {
int row_p = index[ratio*ib + ir];
int row_c = index[ratio*ib + ir + idx->ne[0]/2];
const float * score_p = (const float *)((const char *)score->data + row_p*score->nb[1]) + first_i0;
const float * state_p = (const float *)((const char *)state->data + row_p*state->nb[1]) + first_i0;
const float * score_c = (const float *)((const char *)score->data + row_c*score->nb[1]) + first_i0 + ne0;
const float * state_c = (const float *)((const char *)state->data + row_c*state->nb[1]) + first_i0 + ne0;
for (int j = 0; j < 32; ++j) {
score_l[64*ir + j] = score_p[j];
state_l[64*ir + j] = state_p[j];
}
for (int j = 0; j < 32; ++j) {
score_l[64*ir + 32 + j] = score_c[j];
state_l[64*ir + 32 + j] = state_c[j];
}
}
for (int j = 0; j < 32; ++j) max_l[j] = score_l[j];
for (int ir = 1; ir < 2*ratio; ++ir) {
for (int j = 0; j < 32; ++j) max_l[j] = MAX(max_l[j], score_l[32*ir + j]);
}
for (int j = 0; j < 32; ++j) {
float w = expf(score_l[j] - max_l[j]);
sum_l[j] = w;
res_l[j] = w * state_l[j];
}
for (int ir = 1; ir < 2*ratio; ++ir) {
for (int j = 0; j < 32; ++j) {
float w = expf(score_l[32*ir + j] - max_l[j]);
sum_l[j] += w;
res_l[j] += w * state_l[32*ir + j];
}
}
float * y = (float *)((char *)dst->data + ib*dst->nb[1]) + first_i0;
for (int j = 0; j < 32; ++j) {
y[j] = sum_l[j] > 0 ? res_l[j] / sum_l[j] : 0.0f;
}
}
}
static void ggml_compute_forward_ds4_comp_type1(const struct ggml_compute_params * params,
struct ggml_tensor * dst) {
struct ggml_tensor * state = dst->src[0];
struct ggml_tensor * score = dst->src[1];
struct ggml_tensor * idx = dst->src[2];
GGML_ASSERT(state->type == GGML_TYPE_F32);
GGML_ASSERT(score->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_are_same_shape(score, state));
GGML_ASSERT(state->ne[2] == 1 && state->ne[3] == 1);
GGML_ASSERT(state->ne[0] % 32 == 0);
GGML_ASSERT(dst->ne[0] == state->ne[0]);
GGML_ASSERT( idx->type == GGML_TYPE_I32);
GGML_ASSERT(ggml_nrows(idx) == 1);
int nblock = dst->ne[1];
int ratio = idx->ne[0] / nblock;
GGML_ASSERT(idx->ne[0] % ratio == 0);
int ith = params->ith;
int nth = params->nth;
int ne0 = dst->ne[0];
int n16 = ne0/16;
int nchunk = n16*nblock;
int npt = (nchunk + nth - 1)/nth;
int first = ith*npt;
int last = MIN(first + npt, nchunk);
const int * index = (const int *)idx->data;
// This seems very slightly better than the commented out version below
//
float max_l[16], sum_l[16], res_l[16];
for (int ic = first; ic < last; ++ic) {
int ib = ic / n16;
int i16 = ic - ib*n16;
int first_i0 = 16*i16;
for (int j = 0; j < 16; ++j) {
max_l[j] = -INFINITY;
sum_l[j] = 0;
res_l[j] = 0;
}
for (int ir = 0; ir < ratio; ++ir) {
int row = index[ratio*ib + ir];
const float * score_r = (const float *)((const char *)score->data + row*score->nb[1]) + first_i0;
for (int j = 0; j < 16; ++j) {
float v = score_r[j];
max_l[j] = MAX(max_l[j], v);
}
}
for (int ir = 0; ir < ratio; ++ir) {
int row = index[ratio*ib + ir];
const float * score_r = (const float *)((const char *)score->data + row*score->nb[1]) + first_i0;
const float * state_r = (const float *)((const char *)state->data + row*state->nb[1]) + first_i0;
for (int j = 0; j < 16; ++j) {
float w = expf(score_r[j] - max_l[j]);
sum_l[j] += w;
res_l[j] += w*state_r[j];
}
}
float * y = (float *)((char *)dst->data + ib*dst->nb[1]) + first_i0;
for (int j = 0; j < 16; ++j) {
y[j] = res_l[j] / sum_l[j];
}
}
//size_t work_size = 16*(2*ratio + 3)*sizeof(float);
//GGML_ASSERT(nth*work_size <= params->wsize);
//float * work = (float *)((char *)params->wdata + ith*work_size);
//float * max_l = work;
//float * sum_l = max_l + 16;
//float * res_l = sum_l + 16;
//float * score_l = res_l + 16;
//float * state_l = score_l + 16*ratio;
//for (int ic = first; ic < last; ++ic) {
// int ib = ic / n16;
// int i16 = ic - ib*n16;
// int first_i0 = 16*i16;
// for (int ir = 0; ir < ratio; ++ir) {
// int row = index[ratio*ib + ir];
// const float * score_r = (const float *)((const char *)score->data + row*score->nb[1]) + first_i0;
// const float * state_r = (const float *)((const char *)state->data + row*state->nb[1]) + first_i0;
// for (int j = 0; j < 16; ++j) {
// score_l[16*ir + j] = score_r[j];
// state_l[16*ir + j] = state_r[j];
// }
// }
// for (int j = 0; j < 16; ++j) max_l[j] = score_l[j];
// for (int ir = 1; ir < ratio; ++ir) {
// for (int j = 0; j < 16; ++j) max_l[j] = MAX(max_l[j], score_l[16*ir + j]);
// }
// for (int j = 0; j < 16; ++j) {
// float w = expf(score_l[j] - max_l[j]);
// sum_l[j] = w;
// res_l[j] = w * state_l[j];
// }
// for (int ir = 1; ir < ratio; ++ir) {
// for (int j = 0; j < 16; ++j) {
// float w = expf(score_l[16*ir + j] - max_l[j]);
// sum_l[j] += w;
// res_l[j] += w * state_l[16*ir + j];
// }
// }
// float * y = (float *)((char *)dst->data + ib*dst->nb[1]) + first_i0;
// for (int j = 0; j < 16; ++j) {
// y[j] = sum_l[j] > 0 ? res_l[j] / sum_l[j] : 0.0f;
// }
//}
}
// ggml_compute_forward_latent_attn
@ -26056,6 +26294,14 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml
{
ggml_compute_forward_latent_attn(params, tensor);
} break;
case GGML_OP_DS4_COMP:
{
if (tensor->op_params[0] == 0) {
ggml_compute_forward_ds4_comp_type0(params, tensor);
} else {
ggml_compute_forward_ds4_comp_type1(params, tensor);
}
} break;
case GGML_OP_INDEXER_TOPK:
{
if (!iqk_indexer_topk(tensor, params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, params->ith, params->nth)) {
@ -27134,6 +27380,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
case GGML_OP_HC_POST:
case GGML_OP_MASK_TO_IDX:
case GGML_OP_LATENT_ATTN:
case GGML_OP_DS4_COMP:
{
GGML_ABORT("fatal error"); // TODO: not implemented
}
@ -27882,6 +28129,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
case GGML_OP_HC_POST:
case GGML_OP_MASK_TO_IDX:
case GGML_OP_LATENT_ATTN:
case GGML_OP_DS4_COMP:
{
n_tasks = n_threads;
} break;

131
src/graphs/build_deepseek4.cpp
View File

@ -795,7 +795,6 @@ static ggml_tensor * build_hca_compressed_kv_from_state(
ggml_tensor * comp_pos,
ggml_tensor * norm,
int64_t n_embd_head,
const llm_build_cb & cb,
int il) {
const int64_t n_embd_head_rope = llm.hparams.n_rot;
const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
@ -804,23 +803,8 @@ static ggml_tensor * build_hca_compressed_kv_from_state(
GGML_ASSERT(n_blocks > 0);
GGML_ASSERT(state_read_idxs != nullptr);
ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs);
cb(kv, "hca_kv", il);
kv = ggml_reshape_3d(ctx0, kv, n_embd_head, llama_context::dsv4_runtime::HCA_RATIO, n_blocks);
llm.cb(kv, "hca_comp_kv_rows", il);
auto comp = ggml_ds4_comp(ctx0, kv_state, score_state, state_read_idxs, llama_context::dsv4_runtime::HCA_RATIO, 1);
ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs);
cb(score, "hca_score", il);
score = ggml_reshape_3d(ctx0, score, n_embd_head, llama_context::dsv4_runtime::HCA_RATIO, n_blocks);
llm.cb(score, "hca_comp_score_rows", il);
ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3));
ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
ggml_tensor * weights = ggml_soft_max(ctx0, scores);
ggml_tensor * comp = ggml_mul(ctx0, values, weights);
comp = ggml_sum_rows(ctx0, comp);
comp = ggml_reshape_3d(ctx0, comp, comp->ne[1], comp->ne[2], comp->ne[3]);
//comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3));
llm.cb(comp, "hca_comp_merge", il);
comp = llm.llm_build_norm(ctx0, comp, llm.hparams, norm, nullptr, LLM_NORM_RMS, llm.cb, il);
@ -854,7 +838,7 @@ static ggml_tensor * build_overlap_compressed_kv_from_state(
int64_t ratio,
int64_t n_embd_head,
int il,
const char * tag, const llm_build_cb & cb) {
const char * tag) {
const int64_t n_embd_head_rope = llm.hparams.n_rot;
const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
const int64_t n_blocks = comp_pos ? comp_pos->ne[0] : 0;
@ -862,58 +846,13 @@ static ggml_tensor * build_overlap_compressed_kv_from_state(
GGML_ASSERT(n_blocks > 0);
GGML_ASSERT(state_read_idxs != nullptr);
// Why do we need this?
// 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);
auto kv_state_prev = ggml_view_4d(ctx0, kv_state, n_embd_head, kv_state->ne[1], kv_state->ne[2], kv_state->ne[3],
kv_state->nb[1], kv_state->nb[2], kv_state->nb[3], 0);
auto kv_state_cur = ggml_view_4d(ctx0, kv_state, n_embd_head, kv_state->ne[1], kv_state->ne[2], kv_state->ne[3],
kv_state->nb[1], kv_state->nb[2], kv_state->nb[3], ggml_row_size(kv_state->type, n_embd_head));
auto score_state_prev = ggml_view_4d(ctx0, score_state, n_embd_head, score_state->ne[1], score_state->ne[2], score_state->ne[3],
score_state->nb[1], score_state->nb[2], score_state->nb[3], 0);
auto score_state_cur = ggml_view_4d(ctx0, score_state, n_embd_head, score_state->ne[1], score_state->ne[2], score_state->ne[3],
score_state->nb[1], score_state->nb[2], score_state->nb[3], ggml_row_size(score_state->type, n_embd_head));
ggml_tensor * comp = ggml_ds4_comp(ctx0, kv_state, score_state, state_read_idxs, ratio, 0);
ggml_tensor * prev_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio * n_blocks, 0);
ggml_tensor * cur_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio * n_blocks, ratio * n_blocks);
//ggml_tensor * kv_prev = ggml_get_rows(ctx0, kv_state, prev_idxs);
//kv_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_prev, n_embd_head, ratio * n_blocks, kv_prev->nb[1], 0));
ggml_tensor * kv_prev = ggml_get_rows(ctx0, kv_state_prev, prev_idxs);
cb(kv_prev, tag, il);
kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks);
//ggml_tensor * score_prev = ggml_get_rows(ctx0, score_state, prev_idxs);
//score_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, score_prev, n_embd_head, ratio * n_blocks, score_prev->nb[1], 0));
ggml_tensor * score_prev = ggml_get_rows(ctx0, score_state_prev, prev_idxs);
cb(score_prev, tag, il);
score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks);
//ggml_tensor * kv_cur = ggml_get_rows(ctx0, kv_state, cur_idxs);
//kv_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_cur, n_embd_head, ratio * n_blocks, kv_cur->nb[1],
// ggml_row_size(kv_cur->type, n_embd_head)));
ggml_tensor * kv_cur = ggml_get_rows(ctx0, kv_state_cur, cur_idxs);
cb(kv_cur, tag, il);
kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks);
//ggml_tensor * score_cur = ggml_get_rows(ctx0, score_state, cur_idxs);
//score_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, score_cur, n_embd_head, ratio * n_blocks, score_cur->nb[1],
// ggml_row_size(score_cur->type, n_embd_head)));
ggml_tensor * score_cur = ggml_get_rows(ctx0, score_state_cur, cur_idxs);
cb(score_cur, tag, il);
score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks);
ggml_tensor * values = dsv4_concat_named(ctx0, kv_prev, kv_cur, 1, "dsv4_comp_values");
ggml_tensor * scores = dsv4_concat_named(ctx0, score_prev, score_cur, 1, "dsv4_comp_scores");
values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3));
scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3));
ggml_tensor * weights = ggml_soft_max(ctx0, scores);
ggml_tensor * comp = ggml_mul(ctx0, values, weights);
comp = ggml_sum_rows(ctx0, comp);
//comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3));
comp = ggml_reshape_3d(ctx0, comp, comp->ne[1], comp->ne[2], comp->ne[3]);
llm.cb(comp, tag, il);
comp = llm.llm_build_norm(ctx0, comp, llm.hparams, norm, nullptr, LLM_NORM_RMS, llm.cb, il);
@ -1060,10 +999,10 @@ static ggml_tensor * dsv4_build_lid_top_k(
indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / std::sqrt(float(n_embd_indexer_head * n_indexer_head)));
ggml_tensor * indexer_k = dsv4_comp_get_k(ctx0,
llm.lctx.dsv4.cache.lid_k[(size_t) il],
llm.lctx.dsv4.cache.lid_k[il],
llm.lctx.dsv4.lid_ctx,
n_embd_indexer_head,
llm.lctx.dsv4.cache.lid_k[(size_t) il]->ne[1]/std::max<uint32_t>(1, llm.lctx.dsv4.cache.n_stream));
llm.lctx.dsv4.cache.lid_k[il]->ne[1]/std::max<uint32_t>(1, llm.lctx.dsv4.cache.n_stream));
llm.cb(indexer_k, "lid_k", il);
const int64_t n_stream = std::max<int64_t>(1, indexer_k->ne[3]);
@ -1168,7 +1107,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
qr = llm_build_norm(ctx0, qr, hparams, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, cb, il);
cb(qr, "qr_norm", il);
const int64_t ratio = hparams.dsv4_compress_ratios[(size_t) il];
const int64_t ratio = hparams.dsv4_compress_ratios[il];
const bool use_compress_rope = ratio != 0;
const float freq_base_l = use_compress_rope ? hparams.dsv4_compress_rope_base : freq_base;
const float freq_scale_l = use_compress_rope ? freq_scale : 1.0f;
@ -1201,6 +1140,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
ggml_tensor * kv = llm_build_lora_mm(lctx, ctx0, model.layers[il].wkv_latent, cur);
cb(kv, "wkv", il);
kv = llm_build_norm(ctx0, kv, hparams, model.layers[il].attn_kv_norm, nullptr, LLM_NORM_RMS, cb, il);
cb(kv, "kv_norm", il);
kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
cb(kv, "kv_norm", il);
@ -1226,7 +1166,6 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
cb(kv, "kv_hadamard", il);
}
}
cb(kv, "dsv4_raw_k_before_write", il);
const float kq_scale = 1.0f / std::sqrt(float(n_embd_head));
ggml_tensor * hca_state_kv = nullptr;
@ -1254,8 +1193,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
ggml_tensor * csa_dep = nullptr;
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[(size_t) il], csa_state_kv, 1, "dsv4_csa_source_kv");
ggml_tensor * csa_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.csa_state_score[(size_t) il], csa_state_score, 1, "dsv4_csa_source_score");
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");
ggml_tensor * csa_comp = build_overlap_compressed_kv_from_state(
ctx0, *this,
csa_source_kv, csa_source_score,
@ -1265,9 +1204,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
llama_context::dsv4_runtime::CSA_RATIO,
n_embd_head,
il,
"csa_state_compress", cb);
"csa_state_compress");
ggml_tensor * csa_comp_2d = ggml_reshape_2d(ctx0, csa_comp, n_embd_head, lctx.dsv4.inputs.csa.state_write_idxs->ne[0]);
ggml_tensor * csa_write = dsv4_comp_cpy_k(ctx0, lctx.dsv4.cache.csa_k[(size_t) il], csa_comp_2d, lctx.dsv4.inputs.csa.state_write_idxs, n_embd_head);
ggml_tensor * csa_write = dsv4_comp_cpy_k(ctx0, lctx.dsv4.cache.csa_k[il], csa_comp_2d, lctx.dsv4.inputs.csa.state_write_idxs, n_embd_head);
ggml_build_forward_expand(gf, csa_write);
cb(csa_write, "dsv4_csa_k_write", il);
csa_dep = csa_comp;
@ -1280,8 +1219,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
cb(csa_persist_kv, "csa_persist_kv", il);
ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, lctx.dsv4.inputs.csa.state_persist_src_idxs);
cb(csa_persist_score, "csa_persist_score", il);
ggml_tensor * csa_state_kv_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.csa_state_kv[(size_t) il], csa_persist_kv, lctx.dsv4.inputs.csa.state_persist_dst_idxs);
ggml_tensor * csa_state_score_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.csa_state_score[(size_t) il], csa_persist_score, lctx.dsv4.inputs.csa.state_persist_dst_idxs);
ggml_tensor * csa_state_kv_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.csa_state_kv[il], csa_persist_kv, lctx.dsv4.inputs.csa.state_persist_dst_idxs);
ggml_tensor * csa_state_score_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.csa_state_score[il], csa_persist_score, lctx.dsv4.inputs.csa.state_persist_dst_idxs);
ggml_build_forward_expand(gf, csa_state_kv_write);
ggml_build_forward_expand(gf, csa_state_score_write);
cb(csa_state_kv_write, "dsv4_csa_k_state_persist", il);
@ -1297,8 +1236,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
ggml_tensor * lid_dep = nullptr;
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[(size_t) il], lid_state_kv, 1, "dsv4_lid_source_kv");
ggml_tensor * lid_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.lid_state_score[(size_t) il], lid_state_score, 1, "dsv4_lid_source_score");
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");
ggml_tensor * lid_comp = build_overlap_compressed_kv_from_state(
ctx0, *this,
lid_source_kv, lid_source_score,
@ -1308,14 +1247,14 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
llama_context::dsv4_runtime::CSA_RATIO,
hparams.indexer_head_size,
il,
"lid_state_compress", cb);
"lid_state_compress");
const int hadamard_block = llama_model::hadamard_size((int) hparams.indexer_head_size);
GGML_ASSERT(hadamard_block > 0);
GGML_ASSERT(lid_comp->ne[0] % hadamard_block == 0);
lid_comp = ggml_hadamard(ctx0, lid_comp, hadamard_block);
cb(lid_comp, "lid_state_compress_hadamard", il);
ggml_tensor * lid_comp_2d = ggml_reshape_2d(ctx0, lid_comp, hparams.indexer_head_size, lctx.dsv4.inputs.lid.state_write_idxs->ne[0]);
ggml_tensor * lid_write = dsv4_comp_cpy_k(ctx0, lctx.dsv4.cache.lid_k[(size_t) il], lid_comp_2d, lctx.dsv4.inputs.lid.state_write_idxs, hparams.indexer_head_size);
ggml_tensor * lid_write = dsv4_comp_cpy_k(ctx0, lctx.dsv4.cache.lid_k[il], lid_comp_2d, lctx.dsv4.inputs.lid.state_write_idxs, hparams.indexer_head_size);
ggml_build_forward_expand(gf, lid_write);
cb(lid_write, "dsv4_lid_k_write", il);
lid_dep = lid_comp;
@ -1328,8 +1267,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
cb(lid_persist_kv, "lid_persist_kv", il);
ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, lctx.dsv4.inputs.lid.state_persist_src_idxs);
cb(lid_persist_score, "lid_persist_score", il);
ggml_tensor * lid_state_kv_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.lid_state_kv[(size_t) il], lid_persist_kv, lctx.dsv4.inputs.lid.state_persist_dst_idxs);
ggml_tensor * lid_state_score_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.lid_state_score[(size_t) il], lid_persist_score, lctx.dsv4.inputs.lid.state_persist_dst_idxs);
ggml_tensor * lid_state_kv_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.lid_state_kv[il], lid_persist_kv, lctx.dsv4.inputs.lid.state_persist_dst_idxs);
ggml_tensor * lid_state_score_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.lid_state_score[il], lid_persist_score, lctx.dsv4.inputs.lid.state_persist_dst_idxs);
ggml_build_forward_expand(gf, lid_state_kv_write);
ggml_build_forward_expand(gf, lid_state_score_write);
cb(lid_state_kv_write, "dsv4_lid_k_state_persist", il);
@ -1339,8 +1278,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
if (ratio == llama_context::dsv4_runtime::HCA_RATIO && hca_state_kv != nullptr && hca_state_score != nullptr) {
ggml_tensor * hca_dep = nullptr;
if (lctx.dsv4.inputs.hca.state_write_idxs != nullptr && lctx.dsv4.hca_plan.state_write_idxs.size() > 0) {
ggml_tensor * hca_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.hca_state_kv[(size_t) il], hca_state_kv, 1, "dsv4_hca_source_kv");
ggml_tensor * hca_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.hca_state_score[(size_t) il], hca_state_score, 1, "dsv4_hca_source_score");
ggml_tensor * hca_source_kv = dsv4_concat_named(ctx0, lctx.dsv4.cache.hca_state_kv[il], hca_state_kv, 1, "dsv4_hca_source_kv");
ggml_tensor * hca_source_score = dsv4_concat_named(ctx0, lctx.dsv4.cache.hca_state_score[il], hca_state_score, 1, "dsv4_hca_source_score");
ggml_tensor * hca_comp = build_hca_compressed_kv_from_state(
ctx0, *this,
hca_source_kv, hca_source_score,
@ -1348,9 +1287,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
lctx.dsv4.inputs.hca.state_write_pos,
model.layers[il].attn_comp_norm,
n_embd_head,
cb, il);
il);
ggml_tensor * hca_comp_2d = ggml_reshape_2d(ctx0, hca_comp, n_embd_head, lctx.dsv4.inputs.hca.state_write_idxs->ne[0]);
ggml_tensor * hca_write = dsv4_comp_cpy_k(ctx0, lctx.dsv4.cache.hca_k[(size_t) il], hca_comp_2d, lctx.dsv4.inputs.hca.state_write_idxs, n_embd_head);
ggml_tensor * hca_write = dsv4_comp_cpy_k(ctx0, lctx.dsv4.cache.hca_k[il], hca_comp_2d, lctx.dsv4.inputs.hca.state_write_idxs, n_embd_head);
ggml_build_forward_expand(gf, hca_write);
cb(hca_write, "dsv4_hca_k_write", il);
hca_dep = hca_comp;
@ -1363,8 +1302,8 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
ggml_tensor * hca_persist_score = ggml_get_rows(ctx0, hca_state_score, lctx.dsv4.inputs.hca.state_persist_src_idxs);
cb(hca_persist_kv, "hca_persist_kv", il);
cb(hca_persist_score, "hca_persist_score", il);
ggml_tensor * hca_state_kv_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.hca_state_kv[(size_t) il], hca_persist_kv, lctx.dsv4.inputs.hca.state_persist_dst_idxs);
ggml_tensor * hca_state_score_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.hca_state_score[(size_t) il], hca_persist_score, lctx.dsv4.inputs.hca.state_persist_dst_idxs);
ggml_tensor * hca_state_kv_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.hca_state_kv[il], hca_persist_kv, lctx.dsv4.inputs.hca.state_persist_dst_idxs);
ggml_tensor * hca_state_score_write = dsv4_comp_state_cpy(ctx0, lctx.dsv4.cache.hca_state_score[il], hca_persist_score, lctx.dsv4.inputs.hca.state_persist_dst_idxs);
ggml_build_forward_expand(gf, hca_state_kv_write);
ggml_build_forward_expand(gf, hca_state_score_write);
cb(hca_state_kv_write, "dsv4_hca_k_state_persist", il);
@ -1381,7 +1320,7 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
if (raw_k_write == nullptr) {
llm_build_kv_store(lctx, ctx0, hparams, cparams, kv_self, gf, kv, nullptr, n_tokens, kv_head, cb, il);
}
if (il < (int64_t) kv_self.v_l.size() && kv_self.v_l[(size_t) il] != nullptr) {
if (il < (int64_t) kv_self.v_l.size() && kv_self.v_l[il] != nullptr) {
llm_build_kv_store(lctx, ctx0, hparams, cparams, kv_self, gf, nullptr, kv, n_tokens, kv_head, cb, il);
}
@ -1419,10 +1358,10 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
lctx.dsv4.lid_plan.n_kv > 0 &&
!cparams.k_cache_hadamard) {
ggml_tensor * csa_k = dsv4_comp_get_k(ctx0,
lctx.dsv4.cache.csa_k[(size_t) il],
lctx.dsv4.cache.csa_k[il],
lctx.dsv4.csa_ctx,
n_embd_head,
lctx.dsv4.cache.csa_k[(size_t) il]->ne[1]/std::max<uint32_t>(1, lctx.dsv4.cache.n_stream));
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,
@ -1458,6 +1397,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
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],
@ -1479,10 +1421,10 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
[](int32_t n_visible) { return n_visible > 0; }) &&
!cparams.k_cache_hadamard) {
ggml_tensor * hca_k = dsv4_comp_get_k(ctx0,
lctx.dsv4.cache.hca_k[(size_t) il],
lctx.dsv4.cache.hca_k[il],
lctx.dsv4.hca_ctx,
n_embd_head,
lctx.dsv4.cache.hca_k[(size_t) il]->ne[1]/std::max<uint32_t>(1, lctx.dsv4.cache.n_stream));
lctx.dsv4.cache.hca_k[il]->ne[1]/std::max<uint32_t>(1, lctx.dsv4.cache.n_stream));
const bool use_fattn = cparams.flash_attn;
ggml_tensor * hca_mask = dsv4_build_raw_mask_view(ctx0, lctx.dsv4.inputs.hca.kq_mask, nullptr,
lctx.dsv4.hca_plan.n_kv, n_tokens, hca_k->ne[3], cb, il);
@ -1507,6 +1449,9 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
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]);