ggml: optimize concat op by replacing per-element memcpy with row-level memcpy (#24575)

* ggml: optimize concat op by replacing per-element memcpy with row-level memcpy

* ggml: fix concat offsets for row-level copies

* ggml: add concat row contiguity asserts

* ggml: move concat block size asserts

* ggml: remove redundant concat asserts

* Update ggml/src/ggml-cpu/ops.cpp

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

---------

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
This commit is contained in:
Kartik Sirohi 2026-08-22 14:00:31 +05:30 committed by GitHub
parent 2c6b141efb
commit 369e1cd614
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1896,7 +1896,6 @@ void ggml_compute_forward_repeat_back(
}
// ggml_compute_forward_concat
static void ggml_compute_forward_concat_any(
const ggml_compute_params * params,
ggml_tensor * dst) {
@ -1904,8 +1903,6 @@ static void ggml_compute_forward_concat_any(
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
const size_t len = ggml_type_size(src0->type);
const int ith = params->ith;
const int nth = params->nth;
@ -1914,31 +1911,38 @@ static void ggml_compute_forward_concat_any(
const int32_t dim = ggml_get_op_params_i32(dst, 0);
GGML_ASSERT(dim >= 0 && dim < 4);
GGML_ASSERT(ggml_is_contiguous_rows(src0));
GGML_ASSERT(ggml_is_contiguous_rows(src1));
int64_t o[4] = {0, 0, 0, 0};
if (dim == 0) {
GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
o[dim] = src0->ne[dim]/ggml_blck_size(src0->type);
} else {
o[dim] = src0->ne[dim];
}
const char * x;
// Region 1: copy rows from src0
for (int i3 = 0; i3 < ne03; i3++) {
for (int i2 = ith; i2 < ne02; i2 += nth) {
for (int i1 = 0; i1 < ne01; i1++) {
const char * x = (const char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03;
char * y = ( char *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
memcpy(y, x, ggml_row_size(src0->type, ne00));
}
}
}
// TODO: smarter multi-theading
for (int i3 = 0; i3 < ne3; i3++) {
for (int i2 = ith; i2 < ne2; i2 += nth) {
for (int i1 = 0; i1 < ne1; i1++) {
for (int i0 = 0; i0 < ne0/ggml_blck_size(dst->type); i0++) {
if (i0 < ne00/ggml_blck_size(src0->type) && i1 < ne01 && i2 < ne02 && i3 < ne03) {
x = (const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03;
} else {
x = (const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13;
}
char * y = (char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3;
memcpy(y, x, len);
}
// Region 2: copy rows from src1, offset into dst by o[]
for (int i3 = 0; i3 < ne13; i3++) {
for (int i2 = ith; i2 < ne12; i2 += nth) {
for (int i1 = 0; i1 < ne11; i1++) {
const char * x = (const char *) src1->data + i1*nb11 + i2*nb12 + i3*nb13;
char * y = ( char *) dst->data + (i1 + o[1])*nb1 + (i2 + o[2])*nb2 + (i3 + o[3])*nb3 + o[0]*nb0;
memcpy(y, x, ggml_row_size(src1->type, ne10));
}
}
}
@ -2078,14 +2082,6 @@ void ggml_compute_forward_concat(
ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
if (ggml_is_quantized(src0->type)) {
GGML_ASSERT(ggml_is_contiguous_rows(src0));
GGML_ASSERT(ggml_is_contiguous_rows(src1));
GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
}
switch (src0->type) {
case GGML_TYPE_F16: