From bd342d624fe1ef2165a8db2fcb77cbcb1ebd7dbc Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sat, 25 Jul 2026 08:52:38 +0300 Subject: [PATCH] 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 --- ggml/include/ggml.h | 9 ++ ggml/src/ggml-cuda.cu | 5 + ggml/src/ggml-cuda/concat.cu | 4 + ggml/src/ggml-cuda/ds4_comp.cu | 164 +++++++++++++++++++++ ggml/src/ggml-cuda/ds4_comp.cuh | 3 + ggml/src/ggml.c | 252 +++++++++++++++++++++++++++++++- src/graphs/build_deepseek4.cpp | 131 +++++------------ 7 files changed, 473 insertions(+), 95 deletions(-) create mode 100644 ggml/src/ggml-cuda/ds4_comp.cu create mode 100644 ggml/src/ggml-cuda/ds4_comp.cuh diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index afe4f486..fe9d09ca 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -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 diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index 015277ad..11187faf 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -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 #include @@ -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: diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu index 979da3f9..c8695f2c 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu @@ -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(); diff --git a/ggml/src/ggml-cuda/ds4_comp.cu b/ggml/src/ggml-cuda/ds4_comp.cu new file mode 100644 index 00000000..0d05fdb7 --- /dev/null +++ b/ggml/src/ggml-cuda/ds4_comp.cu @@ -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<<>>(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<<>>(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<<>>(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); + } + +} diff --git a/ggml/src/ggml-cuda/ds4_comp.cuh b/ggml/src/ggml-cuda/ds4_comp.cuh new file mode 100644 index 00000000..be6fd4e1 --- /dev/null +++ b/ggml/src/ggml-cuda/ds4_comp.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_ds4_comp(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index fe82fa46..2d443c94 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -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; diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index 9c88fa54..42e36338 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -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(1, llm.lctx.dsv4.cache.n_stream)); + llm.lctx.dsv4.cache.lid_k[il]->ne[1]/std::max(1, llm.lctx.dsv4.cache.n_stream)); llm.cb(indexer_k, "lid_k", il); const int64_t n_stream = std::max(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(1, lctx.dsv4.cache.n_stream)); + 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, @@ -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(1, lctx.dsv4.cache.n_stream)); + lctx.dsv4.cache.hca_k[il]->ne[1]/std::max(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]);