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

Code before:
                const int64_t d_inner      = ne01;
                const int64_t n_head       = ne02;
                const int64_t n_group      = ne41;
                const int64_t n_seq_tokens = ne12;
                const int64_t n_seqs       = ne13;

                id<MTLComputePipelineState> pipeline = nil;

                if (ne30 == 1) {
                    // Mamba-2
                    pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_SSM_SCAN_F32_GROUP].pipeline;
                } else {
                    pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_SSM_SCAN_F32].pipeline;
                }

                ggml_metal_kargs_ssm_scan args = {
                    /*.d_state      =*/ d_state,
                    /*.d_inner      =*/ d_inner,
                    /*.n_head       =*/ n_head,
                    /*.n_group      =*/ n_group,
                    /*.n_seq_tokens =*/ n_seq_tokens,
                    /*.n_seqs       =*/ n_seqs,
                    /*.s_off        =*/ ggml_nelements(src1) * sizeof(float),
                    /*.nb01         =*/ nb01,
                    /*.nb02         =*/ nb02,
                    /*.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];


Review comment (suggestion):
Should this be simply?

```suggestion
                    /*.s_off        =*/ ggml_nbytes(src1),
```

Code after:
                const uint64_t nb60 = src6->nb[0]; GGML_UNUSED(nb60);

                const int64_t d_state      = ne00;
                const int64_t d_inner      = ne01;
                const int64_t n_head       = ne02;
                const int64_t n_group      = ne41;
                const int64_t n_seq_tokens = ne12;
                const int64_t n_seqs       = ne13;

                id<MTLComputePipelineState> pipeline = nil;

                if (ne30 == 1) {
                    // Mamba-2
                    pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_SSM_SCAN_F32_GROUP].pipeline;
                } else {
                    pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_SSM_SCAN_F32].pipeline;
                }

                ggml_metal_kargs_ssm_scan args = {
                    /*.d_state      =*/ d_state,
                    /*.d_inner      =*/ d_inner,
                    /*.n_head       =*/ n_head,
                    /*.n_group      =*/ n_group,
                    /*.n_seq_tokens =*/ n_seq_tokens,
                    /*.n_seqs       =*/ n_seqs,
                    /*.s_off        =*/ ggml_nelements(src1) * sizeof(float),
                    /*.nb01         =*/ nb01,
                    /*.nb02         =*/ nb02,
                    /*.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];