change to strides

ggml-org · Nov 19, 2024 · 093f087 · 093f087
1 parent 630d65d
commit 093f087
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/ggml-metal/ggml-metal.m b/src/ggml-metal/ggml-metal.m
@@ -3483,7 +3483,7 @@ static void ggml_metal_encode_node(
                 [encoder setComputePipelineState:pipeline];
                 [encoder setBuffer:id_src0 offset:offs_src0        atIndex:0];
                 [encoder setBuffer:id_dst  offset:offs_dst         atIndex:1];
-                [encoder setBytes:&ne00    length:sizeof( int64_t) atIndex:2];
+                [encoder setBytes:&nb01    length:sizeof( int64_t) atIndex:2];
 
                 [encoder dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
             } break;

diff --git a/src/ggml-metal/ggml-metal.metal b/src/ggml-metal/ggml-metal.metal
@@ -1347,9 +1347,9 @@ kernel void kernel_ssm_scan_f32(
 kernel void kernel_argmax(
         device   const void * x,
         device      int32_t * dst,
-        constant    int64_t & ncols,
+        constant    int64_t & nb01,
         uint tgpig[[threadgroup_position_in_grid]]) {
-    device const float * x_row = (device const float *) ((device const char *) x + tgpig * ncols * sizeof(float));
+    device const float * x_row = (device const float *) ((device const char *) x + tgpig * nb01);
 
     dst[tpitg] = 0;