whisper.rn 0.5.1 → 0.5.3

package/cpp/ggml-metal/ggml-metal-ops.h

@@ -39,6 +39,8 @@ size_t wsp_ggml_metal_op_mul_mat_id_extra_ids(const struct wsp_ggml_tensor * op)
 // return true if we should use the FA vector kernel for this op
 bool wsp_ggml_metal_op_flash_attn_ext_use_vec(const struct wsp_ggml_tensor * op);
 
+size_t wsp_ggml_metal_op_flash_attn_ext_extra_pad(const struct wsp_ggml_tensor * op);
+size_t wsp_ggml_metal_op_flash_attn_ext_extra_blk(const struct wsp_ggml_tensor * op);
 size_t wsp_ggml_metal_op_flash_attn_ext_extra_tmp(const struct wsp_ggml_tensor * op);
 
 int wsp_ggml_metal_op_concat            (wsp_ggml_metal_op_t ctx, int idx);
@@ -48,7 +50,9 @@ int wsp_ggml_metal_op_scale             (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_clamp             (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_unary             (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_glu               (wsp_ggml_metal_op_t ctx, int idx);
+int wsp_ggml_metal_op_sum               (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_sum_rows          (wsp_ggml_metal_op_t ctx, int idx);
+int wsp_ggml_metal_op_cumsum            (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_get_rows          (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_set_rows          (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_soft_max          (wsp_ggml_metal_op_t ctx, int idx);
@@ -67,7 +71,9 @@ int wsp_ggml_metal_op_group_norm        (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_norm              (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_rope              (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_im2col            (wsp_ggml_metal_op_t ctx, int idx);
+int wsp_ggml_metal_op_conv_2d           (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_conv_transpose_1d (wsp_ggml_metal_op_t ctx, int idx);
+int wsp_ggml_metal_op_conv_transpose_2d (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_upscale           (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_pad               (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_pad_reflect_1d    (wsp_ggml_metal_op_t ctx, int idx);
@@ -76,6 +82,8 @@ int wsp_ggml_metal_op_timestep_embedding(wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_argmax            (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_argsort           (wsp_ggml_metal_op_t ctx, int idx);
 int wsp_ggml_metal_op_leaky_relu        (wsp_ggml_metal_op_t ctx, int idx);
+int wsp_ggml_metal_op_opt_step_adamw    (wsp_ggml_metal_op_t ctx, int idx);
+int wsp_ggml_metal_op_opt_step_sgd      (wsp_ggml_metal_op_t ctx, int idx);
 
 #ifdef __cplusplus
 }

package/cpp/ggml-metal/ggml-metal.cpp

@@ -193,9 +193,14 @@ static size_t wsp_ggml_backend_metal_buffer_type_get_alloc_size(wsp_ggml_backend
             } break;
         case WSP_GGML_OP_FLASH_ATTN_EXT:
             {
-                if (wsp_ggml_metal_op_flash_attn_ext_use_vec(tensor)) {
-                    res += wsp_ggml_metal_op_flash_attn_ext_extra_tmp(tensor);
-                }
+                res += wsp_ggml_metal_op_flash_attn_ext_extra_pad(tensor);
+                res += wsp_ggml_metal_op_flash_attn_ext_extra_blk(tensor);
+                res += wsp_ggml_metal_op_flash_attn_ext_extra_tmp(tensor);
+            } break;
+        case WSP_GGML_OP_CUMSUM:
+        case WSP_GGML_OP_ARGSORT:
+            {
+                res *= 2;
             } break;
         default:
             break;