Repository: ggerganov/llama.cpp
File: ggml/src/ggml-metal/ggml-metal.m (Objective-C)

Code before:
                    /*.nb13         =*/ nb13,
                    /*.nb21         =*/ nb21,
                    /*.nb22         =*/ nb22,
                    /*.nb31         =*/ nb31,
                    /*.nb41         =*/ nb41,
                    /*.nb42         =*/ nb42,
                    /*.nb43         =*/ nb43,
                    /*.nb51         =*/ nb51,
                    /*.nb52         =*/ nb52,
                    /*.nb53         =*/ nb53,
                };

                [encoder setComputePipelineState:pipeline];
                [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
                [encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
                [encoder setBuffer:id_src2 offset:offs_src2 atIndex:2];
                [encoder setBuffer:id_src3 offset:offs_src3 atIndex:3];
                [encoder setBuffer:id_src4 offset:offs_src4 atIndex:4];
                [encoder setBuffer:id_src5 offset:offs_src5 atIndex:5];
                [encoder setBuffer:id_src6 offset:offs_src6 atIndex:6];
                [encoder setBuffer:id_dst  offset:offs_dst  atIndex:7];
                [encoder setBytes:&args    length:sizeof(args) atIndex:8];

                if (ne30 == 1) {
                    // Mamba-2
                    [encoder setThreadgroupMemoryLength:32*sizeof(float) atIndex:0]; // SIMD size
                    [encoder dispatchThreadgroups:MTLSizeMake(d_inner, n_head, n_seqs) threadsPerThreadgroup:MTLSizeMake(d_state, 1, 1)];
                } else {
                    GGML_ASSERT(d_inner == 1);
                    [encoder dispatchThreadgroups:MTLSizeMake(n_head, n_seqs, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
                }
            } break;
        case GGML_OP_RWKV_WKV6:
            {
                const int64_t B = dst->src[5]->ne[1];
                const int64_t T = dst->src[0]->ne[2];
                const int64_t C = dst->ne[0];
                const int64_t H = dst->src[0]->ne[1];

                GGML_ASSERT(dst->src[5]->type == GGML_TYPE_F32);
                GGML_ASSERT(C % H == 0);
                GGML_ASSERT(C / H == 64);

                size_t offs_src3 = 0;
                size_t offs_src4 = 0;
                size_t offs_src5 = 0;

                id<MTLBuffer> id_src3 = dst->src[3] ? ggml_metal_get_buffer(dst->src[3], &offs_src3) : nil;
                id<MTLBuffer> id_src4 = dst->src[4] ? ggml_metal_get_buffer(dst->src[4], &offs_src4) : nil;
                id<MTLBuffer> id_src5 = dst->src[5] ? ggml_metal_get_buffer(dst->src[5], &offs_src5) : nil;


Review comment (bug):
I think this might be wrong. I'm pretty sure it should actually be sized to fit one sum for each simd group in the threadgroup, so it should be `d_state / <simd_size>`

Code after:
                [encoder setBuffer:id_dst offset:offs_dst       atIndex:2];
                [encoder setBytes:&args   length:sizeof(args)   atIndex:3];

                [encoder setThreadgroupMemoryLength:32*sizeof(float) atIndex:0];

                [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
            } break;
        case GGML_OP_DIAG_MASK_INF:
            {
                const int n_past = ((const int32_t *)(dst->op_params))[0];

                id<MTLComputePipelineState> pipeline = nil;

                if (ne00%8 == 0) {
                    pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_DIAG_MASK_INF_8].pipeline;
                } else {
                    pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_DIAG_MASK_INF].pipeline;
                }

                ggml_metal_kargs_diag_mask_inf args = {
                    /*.ne00 =*/ ne00,
                    /*.ne01 =*/ ne01,
                    /*.n_past =*/ n_past,
                };

                [encoder setComputePipelineState:pipeline];
                [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
                [encoder setBuffer:id_dst  offset:offs_dst  atIndex:1];
                [encoder setBytes:&args  length:sizeof(args) atIndex:2];

                if (ne00%8 == 0) {
                    [encoder dispatchThreadgroups:MTLSizeMake(ne00*ne01*ne02/8, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
                }
                else {
                    [encoder dispatchThreadgroups:MTLSizeMake(ne00, ne01, ne02) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
                }
            } break;
        case GGML_OP_SSM_CONV:
            {
                GGML_ASSERT(src0t == GGML_TYPE_F32);
                GGML_ASSERT(src1t == GGML_TYPE_F32);

                GGML_ASSERT(ggml_is_contiguous(src0));
                GGML_ASSERT(ggml_is_contiguous(src1));

                id<MTLComputePipelineState> pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_SSM_CONV_F32].pipeline;

                ggml_metal_kargs_ssm_conv args = {
                    /*.ne00 =*/ ne00,
                    /*.ne01 =*/ ne01,
                    /*.ne02 =*/ ne02,