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];

                // One shared memory bucket for each simd group in the threadgroup
                if (d_state >= 32) {
                    const int64_t shmem_size = d_state / 32;

                    // The final simd_sum won't work if the number of simd groups is
                    // larger than the size of a single simd group. If this case is
                    // hit at some point, the logic in the second simd_sum could be
                    // expanded to handle this with one more sequential simd_sum to
                    // collapse simd group sums another time.
                    GGML_ASSERT(shmem_size <= 32);

                    // One thread pre element in d_state
                    GGML_ASSERT(d_state <= (int64_t)pipeline.maxTotalThreadsPerThreadgroup);

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

                if (ne30 == 1) {
                    // Mamba-2
                    [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(d_state, 1, 1)];
                }
            } break;
        case GGML_OP_RWKV_WKV6:
            {
                const int64_t B = dst->src[5]->ne[1];

Review comment (suggestion):
You can always allocate `32*sizeof(float)` bytes of shared memory for simplicity:

```suggestion
                    const int64_t shmem_size = 32;
```

Note that shared memory buffers in metal require to have size multiple of 16:

https://github.com/ggml-org/llama.cpp/blob/e2b7621e7c265a6739225125cf9c534f471b3472/ggml/src/ggml-metal/ggml-metal.m#L4912-L4916

By forcing `shmem_size = 32;` you handle this requirement.

Code after:
                    /*.nb03         =*/ nb03,
                    /*.nb11         =*/ nb11,
                    /*.nb12         =*/ nb12,
                    /*.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];

                // One shared memory bucket for each simd group in the threadgroup
                // NOTE: Metal kernels require the buffer size to be multiple of 16 bytes
                //  https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/1443142-setthreadgroupmemorylength
                if (d_state >= 32) {
                    GGML_ASSERT((int64_t)(d_state / 32) <= 32);
                    const int64_t shmem_size = 32;
                    GGML_ASSERT(d_state <= (int64_t)pipeline.maxTotalThreadsPerThreadgroup);
                    [encoder setThreadgroupMemoryLength:(shmem_size)*sizeof(float) atIndex:0];
                }

                if (ne30 == 1) {
                    // Mamba-2
                    [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(d_state, 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];