From ab0b3bd3c846839e9014c03b071f0f550a57ca84 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 31 Aug 2026 23:16:04 +0300 Subject: [PATCH] metal : add concat support for quantized types (#28116) Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731 --- ggml/src/ggml-metal/ggml-metal-device.m | 6 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 24 ++++++++++-- ggml/src/ggml-metal/kernels/quantize.metal | 45 ++++++++++++++++++++++ 3 files changed, 71 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 83344539e..0d8484d00 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1538,6 +1538,12 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return true; case GGML_TYPE_BF16: return has_bfloat; + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + return true; default: return false; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 30ea4ec27..bc8b3c8d4 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -552,8 +552,24 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { const int32_t dim = ((const int32_t *) op->op_params)[0]; + const bool is_q = ggml_is_quantized(op->type); + + // for quantized types, concat is done at the block level (nb0 == type_size == block size) + int32_t ne00_arg = ne00; + int32_t ne10_arg = ne10; + int32_t ne0_arg = ne0; + if (is_q) { + const int32_t blck = ggml_blck_size(op->type); + GGML_ASSERT(ne00 % blck == 0); + GGML_ASSERT(ne10 % blck == 0); + GGML_ASSERT(ne0 % blck == 0); + ne00_arg = ne00/blck; + ne10_arg = ne10/blck; + ne0_arg = ne0/blck; + } + ggml_metal_kargs_concat args = { - /*.ne00 =*/ ne00, + /*.ne00 =*/ ne00_arg, /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, /*.ne03 =*/ ne03, @@ -561,7 +577,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { /*.nb01 =*/ nb01, /*.nb02 =*/ nb02, /*.nb03 =*/ nb03, - /*.ne10 =*/ ne10, + /*.ne10 =*/ ne10_arg, /*.ne11 =*/ ne11, /*.ne12 =*/ ne12, /*.ne13 =*/ ne13, @@ -569,7 +585,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { /*.nb11 =*/ nb11, /*.nb12 =*/ nb12, /*.nb13 =*/ nb13, - /*.ne0 =*/ ne0, + /*.ne0 =*/ ne0_arg, /*.ne1 =*/ ne1, /*.ne2 =*/ ne2, /*.ne3 =*/ ne3, @@ -588,7 +604,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); - int nth = std::min(256, ne0); + int nth = std::min(256, ne0_arg); // when rows are small, we can batch them together in a single threadgroup int nrptg = 1; diff --git a/ggml/src/ggml-metal/kernels/quantize.metal b/ggml/src/ggml-metal/kernels/quantize.metal index 59d0afe96..42ca6d74a 100644 --- a/ggml/src/ggml-metal/kernels/quantize.metal +++ b/ggml/src/ggml-metal/kernels/quantize.metal @@ -207,6 +207,51 @@ template [[host_name("kernel_concat_i16")]] kernel kernel_concat_t kernel_conca template [[host_name("kernel_concat_i32")]] kernel kernel_concat_t kernel_concat; template [[host_name("kernel_concat_i64")]] kernel kernel_concat_t kernel_concat; +template +kernel void kernel_concat_q( + constant ggml_metal_kargs_concat & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + + // note: for quantized types, the args are in units of blocks (nb0 == type_size) + const int i3 = tgpig.z; + const int i2 = tgpig.y; + const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y; + + if (i1 >= args.ne1) { + return; + } + + int o[4] = {0, 0, 0, 0}; + o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03)); + + for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) { + device const block_q * x; + + if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) { + x = (device const block_q *)(src0 + (i3 )*args.nb03 + (i2 )*args.nb02 + (i1 )*args.nb01 + (i0 )*args.nb00); + } else { + x = (device const block_q *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10); + } + + device block_q * y = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); + + *y = *x; + } +} + +typedef decltype(kernel_concat_q) kernel_concat_q_t; + +template [[host_name("kernel_concat_q4_0")]] kernel kernel_concat_q_t kernel_concat_q; +template [[host_name("kernel_concat_q4_1")]] kernel kernel_concat_q_t kernel_concat_q; +template [[host_name("kernel_concat_q5_0")]] kernel kernel_concat_q_t kernel_concat_q; +template [[host_name("kernel_concat_q5_1")]] kernel kernel_concat_q_t kernel_concat_q; +template [[host_name("kernel_concat_q8_0")]] kernel kernel_concat_q_t kernel_concat_q; + template kernel void kernel_get_rows_q( constant ggml_metal_kargs_get_rows & args,