diff --git a/ggml/src/ggml-metal/kernels/gated_delta_net.metal b/ggml/src/ggml-metal/kernels/gated_delta_net.metal index 8ff8afb8f9..5e4861ece3 100644 --- a/ggml/src/ggml-metal/kernels/gated_delta_net.metal +++ b/ggml/src/ggml-metal/kernels/gated_delta_net.metal @@ -15,7 +15,7 @@ kernel void kernel_gated_delta_net_impl( device const char * b, device const char * s, device char * dst, - device char * state_out, + device char * dst_fuse, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { @@ -69,7 +69,7 @@ kernel void kernel_gated_delta_net_impl( // when fused with the cache cpy, write the snapshots straight into the cache buffer using // the slot stride; otherwise append them after the attn scores (nb_out == 0) const bool fused = args.nb_out > 0; - const device float * state_out_ = fused ? (device float *)state_out : (device float *)dst + attn_size; + const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + attn_size; const uint slot_stride = fused ? (uint)args.nb_out : state_size_per_snap; for (short t = 0; t < args.ne22; t++) { @@ -123,7 +123,7 @@ kernel void kernel_gated_delta_net_impl( if (K > 1) { const int target_slot = (int)args.ne22 - 1 - (int)t; if (target_slot >= 0 && target_slot < (int)K) { - device float * dst_state = (device float *)state_out_ + (uint)target_slot * slot_stride + state_out_base; + device float * dst_state = (device float *)state_out + (uint)target_slot * slot_stride + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; @@ -133,7 +133,7 @@ kernel void kernel_gated_delta_net_impl( } if (K == 1) { - device float * dst_state = (device float *)state_out_ + state_out_base; + device float * dst_state = (device float *)state_out + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; @@ -165,7 +165,7 @@ kernel void kernel_gated_delta_net_impl( device const char * b, device const char * s, device char * dst, - device char * state_out, + device char * dst_fuse, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { @@ -241,10 +241,10 @@ kernel void kernel_gated_delta_net_impl( // when fused with the cache cpy, write the snapshots straight into the cache buffer using // the slot stride; otherwise append them after the attn scores (nb_out == 0) const bool fused = args.nb_out > 0; - const device float * state_out_ = fused ? (device float *)state_out : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v; + const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v; const uint slot_stride = fused ? (uint)args.nb_out : S_v*S_v; - device float * dst_state = (device float *)state_out_ + (i23*args.ne21 + i21)*slot_stride + i20; + device float * dst_state = (device float *)state_out + (i23*args.ne21 + i21)*slot_stride + i20; device T * dstt_state = (device T *) (dst_state); FOR_UNROLL (short j = 0; j < NSG; j++) {