cuda : concat implementation for quantized types (#25303)

* cuda : concat implementation for quantized types

* chore : apply am17an clever suggestion to shorten the code

---------

Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
This commit is contained in:
fairydreaming
2026-07-05 23:26:24 +08:00
committed by GitHub
co-authored by Stanisław Szymczyk
parent a4107133a6
commit 78d2f52468
2 changed files with 50 additions and 26 deletions
+15 -3
View File
@@ -152,8 +152,8 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml
src0_d + i3*(src0->nb[3] / sizeof(T)), src0_d + i3*(src0->nb[3] / sizeof(T)),
src1_d + i3*(src1->nb[3] / sizeof(T)), src1_d + i3*(src1->nb[3] / sizeof(T)),
dst_d + i3*( dst->nb[3] / sizeof(T)), dst_d + i3*( dst->nb[3] / sizeof(T)),
src0->ne[0], src0->ne[1], src0->ne[2], ggml_row_size(src0->type, src0->ne[0])/sizeof(T), src0->ne[1], src0->ne[2],
dst->ne[0], dst->ne[1], dst->ne[2], dim, stream); ggml_row_size(dst->type, dst->ne[0])/sizeof(T), dst->ne[1], dst->ne[2], dim, stream);
} }
} else { } else {
const size_t size0 = ggml_nbytes(src0); const size_t size0 = ggml_nbytes(src0);
@@ -163,6 +163,8 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml
CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream)); CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream));
} }
} else { } else {
GGML_ASSERT(!ggml_is_quantized(src0->type));
dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]);
auto launch_kernel = [&](auto dim) { auto launch_kernel = [&](auto dim) {
concat_non_cont<T, dim><<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>( concat_non_cont<T, dim><<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
@@ -204,7 +206,16 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
GGML_ASSERT(src0->type == src1->type); GGML_ASSERT(src0->type == src1->type);
GGML_ASSERT(dst->type == src0->type); GGML_ASSERT(dst->type == src0->type);
GGML_ASSERT(!ggml_is_quantized(src0->type));
if (ggml_is_quantized(src0->type)) {
GGML_ASSERT(ggml_is_contiguous(src0));
GGML_ASSERT(ggml_is_contiguous(src1));
GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
// if tensors are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors
concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
} else {
GGML_ASSERT(ggml_blck_size(src0->type) == 1); GGML_ASSERT(ggml_blck_size(src0->type) == 1);
switch (ggml_type_size(src0->type)) { switch (ggml_type_size(src0->type)) {
@@ -224,4 +235,5 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type)); GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type));
break; break;
} }
}
} }
+14 -2
View File
@@ -5387,12 +5387,24 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
ggml_type src1_type = op->src[1]->type; ggml_type src1_type = op->src[1]->type;
return src0_type == src1_type && return src0_type == src1_type &&
src0_type == op->type && src0_type == op->type &&
(
(
ggml_is_quantized(src0_type) &&
ggml_is_contiguous(op->src[0]) &&
ggml_is_contiguous(op->src[1]) &&
op->src[0]->ne[0] % ggml_blck_size(src0_type) == 0 &&
op->src[1]->ne[0] % ggml_blck_size(src0_type) == 0
) || (
!ggml_is_quantized(src0_type) && !ggml_is_quantized(src0_type) &&
ggml_blck_size(src0_type) == 1 && ggml_blck_size(src0_type) == 1 &&
(ggml_type_size(src0_type) == 1 || (
ggml_type_size(src0_type) == 1 ||
ggml_type_size(src0_type) == 2 || ggml_type_size(src0_type) == 2 ||
ggml_type_size(src0_type) == 4 || ggml_type_size(src0_type) == 4 ||
ggml_type_size(src0_type) == 8); ggml_type_size(src0_type) == 8
)
)
);
} break; } break;
case GGML_OP_CONV_TRANSPOSE_1D: case GGML_OP_CONV_TRANSPOSE_1D:
{ {