diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d16a2d02f..b52ecd2e7 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -1946,6 +1946,13 @@ extern "C" { struct ggml_tensor * b, struct ggml_tensor * c); + GGML_API struct ggml_tensor * ggml_get_rows_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool same_type, + bool dim0); + // a TD [n_embd, ne1, ne2, ne3] // b TS [n_embd, n_rows, ne02, ne03] | ne02 == ne2, ne03 == ne3 // c I64 [n_rows, ne11, ne12, 1] | c[i] in [0, ne1) diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index 11187faf8..ad0db58e2 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -4841,6 +4841,9 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons } break; case GGML_OP_GET_ROWS: { + if (op->type == op->src[0]->type) { + return true; + } switch (op->src[0]->type) { case GGML_TYPE_F16: case GGML_TYPE_F32: diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu index fd67dce1d..40c677c48 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu @@ -139,6 +139,34 @@ static void get_rows_cuda_float(const ggml_tensor * src0, const ggml_tensor * sr GGML_UNUSED(dst); } +template +static __global__ void k_get_rows_dim0(int n, + 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 data_t * __restrict__ src, const int * __restrict__ idx, data_t * __restrict__ dst) { + int i1 = blockIdx.x; + int i2 = blockIdx.y; + int i3 = blockIdx.z; + src += nb01*i1 + nb02*i2 + nb03*i3; + idx += nb11*i1 + nb12*i2 + nb13*i3; + dst += nb1*i1 + nb2*i2 + nb3*i3; + + for (int j = threadIdx.x; j < n; j += blockDim.x) dst[j] = src[idx[j]]; +} + +template +static __global__ void k_get_rows_dim1(int n, + size_t nb01, size_t nb02, size_t nb03, size_t nb11, size_t nb12, size_t nb1, size_t nb2, size_t nb3, + const data_t * __restrict__ src, const int * __restrict__ idx, data_t * __restrict__ dst) { + int i1 = blockIdx.x; + int i2 = blockIdx.y; + int i3 = blockIdx.z; + idx += nb11*i2 + nb12*i3; + src += nb01*idx[i1] + nb02*i2 + nb03*i3; + dst += nb1*i1 + nb2*i2 + nb3*i3; + + for (int j = threadIdx.x; j < n; j += blockDim.x) dst[j] = src[j]; +} + void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -147,8 +175,60 @@ void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { float * dst_d = (float *)dst->data; cudaStream_t stream = ctx.stream(); - GGML_ASSERT(src1->type == GGML_TYPE_I32); + + if (dst->type != GGML_TYPE_F32 || dst->op_params[0] == 1) { + constexpr int k_block_size = 512; + GGML_ASSERT(src0->type == dst->type); + if (dst->op_params[0] == 1) { + //printf("%s(%s, %s) - dim0\n", __func__, src0->name, ggml_type_name(src0->type)); + GGML_ASSERT(src0->ne[1] == src1->ne[1] && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3]); + GGML_ASSERT(src0->ne[0] >= src1->ne[0]); + GGML_ASSERT(dst->ne[0] == src1->ne[0]); + auto type_size = ggml_type_size(dst->type); + GGML_ASSERT(type_size == 2 || type_size == 4); + dim3 grid(dst->ne[1], dst->ne[2], dst->ne[3]); + if (type_size == 2) { + k_get_rows_dim0<<>>(dst->ne[0], + src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, + src1->nb[1]/type_size, src1->nb[2]/type_size, src1->nb[3]/type_size, + dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, + (const uint16_t *)src0->data, (const int *)src1->data, (uint16_t *)dst->data); + } else { + k_get_rows_dim0<<>>(dst->ne[0], + src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, + src1->nb[1]/type_size, src1->nb[2]/type_size, src1->nb[3]/type_size, + dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, + (const uint32_t *)src0->data, (const int *)src1->data, (uint32_t *)dst->data); + } + } else { + //printf("%s(%s, %s) - dim1\n", __func__, src0->name, ggml_type_name(src0->type)); + GGML_ASSERT(src0->ne[2] == src1->ne[1]); + GGML_ASSERT(src0->ne[3] == src1->ne[2]); + GGML_ASSERT(src0->ne[3] == 1); + GGML_ASSERT(dst->ne[0] == src0->ne[0] && dst->ne[1] == src1->ne[0] && dst->ne[2] == src1->ne[1] && dst->ne[2]); + + auto row_size = ggml_row_size(dst->type, dst->ne[0]); + GGML_ASSERT(row_size % 2 == 0); + auto type_size = row_size % 4 == 0 ? 4 : 2; + dim3 grid(dst->ne[1], dst->ne[2], dst->ne[3]); + if (type_size == 2) { + k_get_rows_dim1<<>>(dst->ne[0], + src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, + src1->nb[1]/type_size, src1->nb[2]/type_size, + dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, + (const uint16_t *)src0->data, (const int *)src1->data, (uint16_t *)dst->data); + } else { + k_get_rows_dim1<<>>(dst->ne[0], + src0->nb[1]/type_size, src0->nb[2]/type_size, src0->nb[3]/type_size, + src1->nb[1]/type_size, src1->nb[2]/type_size, + dst->nb[1]/type_size, dst->nb[2]/type_size, dst->nb[3]/type_size, + (const uint32_t *)src0->data, (const int *)src1->data, (uint32_t *)dst->data); + } + } + return; + } + GGML_ASSERT(dst->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32)); GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 94dd995e1..6a097177f 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -8952,6 +8952,45 @@ struct ggml_tensor * ggml_get_rows( return result; } +struct ggml_tensor * ggml_get_rows_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool same_type, + bool dim0) { + if (dim0) { + GGML_ASSERT(same_type); + GGML_ASSERT(!ggml_is_quantized(a->type)); + GGML_ASSERT(ggml_type_size(a->type) % 2 == 0); // we do not support dim0 get_rows on int8, fp8 data + } + if (!dim0 && !same_type) { + return ggml_get_rows(ctx, a, b); + } + GGML_ASSERT(b->type == GGML_TYPE_I32); + if (dim0) { + GGML_ASSERT(a->ne[1] == b->ne[1] && a->ne[2] == b->ne[2] && a->ne[3] == b->ne[3]); + GGML_ASSERT(a->ne[0] >= b->ne[0]); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, b->ne[0], a->ne[1], a->ne[2], a->ne[3]); + result->op = GGML_OP_GET_ROWS; + result->op_params[0] = 1; + result->src[0] = a; + result->src[1] = b; + return result; + } + GGML_ASSERT(a->ne[2] == b->ne[1]); + GGML_ASSERT(a->ne[3] == b->ne[2]); + GGML_ASSERT(b->ne[3] == 1); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); + result->op = GGML_OP_GET_ROWS; + result->src[0] = a; + result->src[1] = b; + return result; + +} + + // ggml_get_rows_back struct ggml_tensor * ggml_get_rows_back( @@ -19572,12 +19611,85 @@ static void ggml_compute_forward_get_rows_f32( } } +static void ggml_compute_forward_get_rows_any( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src1->type == GGML_TYPE_I32); + GGML_ASSERT(src0->type == dst->type); + + const int ith = params->ith; + const int nth = params->nth; + + int nrows = ggml_nrows(dst); + int npt = (nrows + nth - 1)/nth; + int first = npt*ith; + int last = MIN(first + npt, nrows); + + if (dst->op_params[0] == 1) { + size_t type_size = ggml_type_size(dst->type); + GGML_ASSERT(type_size == 2 || type_size == 4); + GGML_ASSERT(src0->ne[1] == src1->ne[1] && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3]); + GGML_ASSERT(src0->ne[0] >= src1->ne[0]); + GGML_ASSERT(dst->ne[0] == src1->ne[0]); + for (int ir = first; ir < last; ++ir) { + int i3 = ir/(dst->ne[1]*dst->ne[2]); + int i2 = (ir - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; + int i1 = ir - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; + const int * idx = (const int *)((const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3]); + const char * cx = (const char *)((const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]); + char * cy = ( char *)(( char *) dst->data + i1* dst->nb[1] + i2* dst->nb[2] + i3* dst->nb[3]); + if (type_size == 2) { + const uint16_t * x = (const uint16_t *)cx; + uint16_t * y = ( uint16_t *)cy; + for (int j = 0; j < (int)dst->ne[0]; ++j) { + y[j] = x[idx[j]]; + } + } else { + const uint32_t * x = (const uint32_t *)cx; + uint32_t * y = ( uint32_t *)cy; + for (int j = 0; j < (int)dst->ne[0]; ++j) { + y[j] = x[idx[j]]; + } + } + } + return; + } + + GGML_ASSERT(src0->ne[2] == src1->ne[1]); + GGML_ASSERT(src0->ne[3] == src1->ne[2]); + GGML_ASSERT(src0->ne[3] == 1); + GGML_ASSERT(dst->ne[0] == src0->ne[0] && dst->ne[1] == src1->ne[0] && dst->ne[2] == src1->ne[1] && dst->ne[2]); + + size_t row_size = ggml_row_size(dst->type, dst->ne[0]); + //GGML_ASSERT(row_size % 2 == 0); + //size_t type_size = row_size % 4 == 0 ? 4 : 2; + + for (int ir = first; ir < last; ++ir) { + int i3 = ir/(dst->ne[1]*dst->ne[2]); + int i2 = (ir - i3*dst->ne[1]*dst->ne[2])/dst->ne[1]; + int i1 = ir - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1]; + const int * idx = (const int *)((const char *)src1->data + i2*src1->nb[1] + i3*src1->nb[2]); + const char * cx = (const char *)src0->data + idx[i1]*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]; + char * cy = ( char *) dst->data + i1 * dst->nb[1] + i2* dst->nb[2] + i3* dst->nb[3]; + memcpy(cy, cx, row_size); + } +} + static void ggml_compute_forward_get_rows( const struct ggml_compute_params * params, struct ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; + if (dst->type != GGML_TYPE_F32) { + ggml_compute_forward_get_rows_any(params, dst); + return; + } + switch (src0->type) { case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index d64357328..d7a8c8cec 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -1118,7 +1118,10 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu ggml_tensor * cache, const auto & extra_ctx, const std::string & tag, int n_swa_eff) { auto n_stream = std::max(1, lctx.dsv4.cache.n_stream); - auto extra_k = dsv4_comp_get_k(ctx0, cache, extra_ctx, n_embd_head, cache->ne[1]/n_stream); + auto extra_k = cache; + if (extra_k->ne[1] > 1) { + extra_k = dsv4_comp_get_k(ctx0, cache, extra_ctx, n_embd_head, cache->ne[1]/n_stream); + } if (cparams.flash_attn) { extra_mask = dsv4_pad_mask_tokens(ctx0, extra_mask, n_tokens); } @@ -1146,7 +1149,6 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu auto attn = dsv4_build_attn(ctx0, hparams, cparams, q, k_all, k_all, kq_mask, model.layers[il].attn_sinks, kq_scale, cb, il, n_swa_eff, gf); return attn; - //return std::make_pair(k_all, kq_mask); }; auto num_streams = [] (const auto & comp) { @@ -1160,16 +1162,24 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu lctx.dsv4.lid_plan.n_kv > 0 && !cparams.k_cache_hadamard) { auto csa_mask = lctx.dsv4.inputs.csa.kq_mask; + auto csa_kv = lctx.dsv4.cache.csa_k[il]; if (hparams.indexer_top_k < lctx.dsv4.inputs.csa.kq_mask->ne[0]) { auto top_k = dsv4_build_lid_top_k(ctx0, llm, 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, num_streams(lctx.dsv4.csa_ctx), cb, il), - top_k); - cb(csa_mask, "csa_mask", il); + if (n_tokens == 1) { + // When we are dealing with a single token, we can just use ggml_get_rows_ext to get the + // selected rows from the CSA cache and setup the corresponding mask. This makes the + // raw_kv and csa_kv concetenation much less expensive for long context. + csa_kv = ggml_get_rows_ext(ctx0, csa_kv, top_k, true, false); + csa_kv = ggml_reshape_3d(ctx0, csa_kv, csa_kv->ne[0], 1, csa_kv->ne[1]); + csa_mask = ggml_get_rows_ext(ctx0, csa_mask, top_k, true, true); + } else { + 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, num_streams(lctx.dsv4.csa_ctx), cb, il), top_k); + cb(csa_mask, "csa_mask", il); + } } int n_csa = hparams.n_swa + hparams.indexer_top_k; - attn = build_the_attn(raw_k, raw_mask, csa_mask, lctx.dsv4.cache.csa_k[il], lctx.dsv4.csa_ctx, "csa", n_csa); + attn = build_the_attn(raw_k, raw_mask, csa_mask, csa_kv, lctx.dsv4.csa_ctx, "csa", n_csa); cb(attn, "attn_csa", il); } else if (ratio == llama_context::dsv4_runtime::HCA_RATIO && lctx.dsv4.inputs.hca.kq_mask != nullptr &&