LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Revert "DS4: faster long-context TG (#2201)"

Browse Source 
This reverts commit 707374b3c3.
This commit is contained in:
Kawrakow
2026-08-04 07:02:51 +00:00
parent c013cd876b
commit 862828d02a
5 changed files with 9 additions and 221 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/include/ggml.h
-7
View File
@@ -1946,13 +1946,6 @@ 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)
ggml/src/ggml-cuda.cu
-3
View File
@@ -4841,9 +4841,6 @@ 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:
ggml/src/ggml-cuda/getrows.cu
+1 -81
View File
@@ -139,34 +139,6 @@ static void get_rows_cuda_float(const ggml_tensor * src0, const ggml_tensor * sr
GGML_UNUSED(dst);
}
template <typename data_t>
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 <typename data_t>
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];
@@ -175,60 +147,8 @@ 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<<<grid, k_block_size, 0, ctx.stream()>>>(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<<<grid, k_block_size, 0, ctx.stream()>>>(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<<<grid, k_block_size, 0, ctx.stream()>>>(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<<<grid, k_block_size, 0, ctx.stream()>>>(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));
ggml/src/ggml.c
-112
View File
@@ -8952,45 +8952,6 @@ 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(
@@ -19736,85 +19697,12 @@ 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:
src/graphs/build_deepseek4.cpp
+8 -18
View File
@@ -1118,10 +1118,7 @@ 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<uint32_t>(1, lctx.dsv4.cache.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);
}
auto 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);
}
@@ -1149,6 +1146,7 @@ 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) {
@@ -1162,24 +1160,16 @@ 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);
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);
}
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, csa_kv, lctx.dsv4.csa_ctx, "csa", n_csa);
attn = build_the_attn(raw_k, raw_mask, csa_mask, lctx.dsv4.cache.csa_k[il], 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 &&