diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 3c81805b844c..79e686fb38e3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -827,6 +827,7 @@ struct vk_device_struct { vk_pipeline pipeline_rwkv_wkv7_f32; // [size_idx][kda] where size_idx: 0=d32, 1=d64, 2=d128 vk_pipeline pipeline_gated_delta_net[3][2]; + vk_pipeline pipeline_gated_delta_net_f16[3][2]; vk_pipeline pipeline_ssm_scan_f32_d128; vk_pipeline pipeline_ssm_scan_f32_d256; vk_pipeline pipeline_ssm_conv_f32; @@ -4589,12 +4590,23 @@ static void ggml_vk_load_shaders(vk_device& device) { {"gated_delta_net_f32_d64", "gated_delta_net_f32_d64_kda"}, {"gated_delta_net_f32_d128", "gated_delta_net_f32_d128_kda"}, }; + const char * gdn_names_f16[][2] = { + {"gated_delta_net_f16_d32", "gated_delta_net_f16_d32_kda"}, + {"gated_delta_net_f16_d64", "gated_delta_net_f16_d64_kda"}, + {"gated_delta_net_f16_d128", "gated_delta_net_f16_d128_kda"}, + }; for (uint32_t si = 0; si < 3; si++) { for (uint32_t kda = 0; kda < 2; kda++) { ggml_vk_create_pipeline(device, device->pipeline_gated_delta_net[si][kda], gdn_names[si][kda], gated_delta_net_f32_len, gated_delta_net_f32_data, "main", 7, sizeof(vk_op_gated_delta_net_push_constants), {1, 1, 1}, {gdn_sizes[si], kda}, 1); + if (device->fp16) { + ggml_vk_create_pipeline(device, device->pipeline_gated_delta_net_f16[si][kda], + gdn_names_f16[si][kda], gated_delta_net_f16_len, gated_delta_net_f16_data, + "main", 7, sizeof(vk_op_gated_delta_net_push_constants), + {1, 1, 1}, {gdn_sizes[si], kda}, 1); + } } } } @@ -9540,6 +9552,9 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const case 128: si = 2; break; default: return nullptr; } + if (ctx->device->fp16 && ctx->device->pipeline_gated_delta_net_f16[si][kda]) { + return ctx->device->pipeline_gated_delta_net_f16[si][kda]; + } return ctx->device->pipeline_gated_delta_net[si][kda]; } return nullptr; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp b/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp index 1fdf889e8246..2c9a70d47a87 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp @@ -1,6 +1,9 @@ #version 450 #extension GL_EXT_control_flow_attributes : require +#ifdef FLOAT16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif layout(constant_id = 0) const uint S_V = 128; layout(constant_id = 1) const uint KDA = 0; @@ -19,17 +22,17 @@ layout(push_constant) uniform Parameters { float scale; }; -layout(binding = 0) readonly buffer QBuf { FLOAT_TYPE data_q[]; }; -layout(binding = 1) readonly buffer KBuf { FLOAT_TYPE data_k[]; }; -layout(binding = 2) readonly buffer VBuf { FLOAT_TYPE data_v[]; }; -layout(binding = 3) readonly buffer GBuf { FLOAT_TYPE data_g[]; }; -layout(binding = 4) readonly buffer BetaBuf { FLOAT_TYPE data_beta[]; }; -layout(binding = 5) readonly buffer StateBuf { FLOAT_TYPE data_state[]; }; -layout(binding = 6) buffer DstBuf { FLOAT_TYPE data_dst[]; }; +layout(binding = 0) readonly buffer QBuf { A_TYPE data_q[]; }; +layout(binding = 1) readonly buffer KBuf { A_TYPE data_k[]; }; +layout(binding = 2) readonly buffer VBuf { A_TYPE data_v[]; }; +layout(binding = 3) readonly buffer GBuf { A_TYPE data_g[]; }; +layout(binding = 4) readonly buffer BetaBuf { A_TYPE data_beta[]; }; +layout(binding = 5) readonly buffer StateBuf { A_TYPE data_state[]; }; +layout(binding = 6) buffer DstBuf { D_TYPE data_dst[]; }; -shared FLOAT_TYPE s_k[S_V]; -shared FLOAT_TYPE s_q[S_V]; -shared FLOAT_TYPE s_g[S_V]; // KDA only: cached exp(g[i]) +shared A_TYPE s_k[S_V]; +shared A_TYPE s_q[S_V]; +shared A_TYPE s_g[S_V]; void main() { const uint head_id = gl_WorkGroupID.x; @@ -54,25 +57,25 @@ void main() { const uint k_off = q_off; const uint v_off = seq_id * sv3 + t * sv2 + head_id * sv1; - s_q[col] = FLOAT_TYPE(data_q[q_off + col]); - s_k[col] = FLOAT_TYPE(data_k[k_off + col]); + s_q[col] = data_q[q_off + col]; + s_k[col] = data_k[k_off + col]; const uint gb_off = seq_id * sb3 + t * sb2 + head_id * sb1; if (KDA != 0) { const uint g_base = gb_off * S_V; - s_g[col] = exp(FLOAT_TYPE(data_g[g_base + col])); + s_g[col] = A_TYPE(exp(float(data_g[g_base + col]))); } barrier(); - const FLOAT_TYPE v_val = FLOAT_TYPE(data_v[v_off + col]); - const FLOAT_TYPE beta_val = FLOAT_TYPE(data_beta[gb_off]); + const float v_val = float(data_v[v_off + col]); + const float beta_val = float(data_beta[gb_off]); if (KDA == 0) { - const FLOAT_TYPE g_val = exp(FLOAT_TYPE(data_g[gb_off])); + const float g_val = exp(float(data_g[gb_off])); - FLOAT_TYPE kv_col = 0.0; + float kv_col = 0.0; [[unroll]] for (uint i = 0; i < S_V; i += 4) { kv_col += dot( vec4(state[i], state[i+1], state[i+2], state[i+3]), @@ -80,21 +83,22 @@ void main() { ); } - FLOAT_TYPE delta_col = (v_val - g_val * kv_col) * beta_val; + float delta_col = (v_val - g_val * kv_col) * beta_val; - FLOAT_TYPE attn_col = 0.0; + float attn_col = 0.0; [[unroll]] for (uint i = 0; i < S_V; i += 4) { vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]); vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]); sv = g_val * sv + kv * delta_col; - state[i] = sv.x; state[i+1] = sv.y; state[i+2] = sv.z; state[i+3] = sv.w; + state[i] = FLOAT_TYPE(sv.x); state[i+1] = FLOAT_TYPE(sv.y); + state[i+2] = FLOAT_TYPE(sv.z); state[i+3] = FLOAT_TYPE(sv.w); attn_col += dot(sv, vec4(s_q[i], s_q[i+1], s_q[i+2], s_q[i+3])); } - data_dst[attn_off + col] = attn_col * scale; + data_dst[attn_off + col] = D_TYPE(attn_col * scale); } else { - FLOAT_TYPE kv_col = 0.0; + float kv_col = 0.0; [[unroll]] for (uint i = 0; i < S_V; i += 4) { vec4 gv = vec4(s_g[i], s_g[i+1], s_g[i+2], s_g[i+3]); vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]); @@ -102,20 +106,21 @@ void main() { kv_col += dot(gv * sv, kv); } - FLOAT_TYPE delta_col = (v_val - kv_col) * beta_val; + float delta_col = (v_val - kv_col) * beta_val; - FLOAT_TYPE attn_col = 0.0; + float attn_col = 0.0; [[unroll]] for (uint i = 0; i < S_V; i += 4) { vec4 gv = vec4(s_g[i], s_g[i+1], s_g[i+2], s_g[i+3]); vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]); vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]); sv = gv * sv + kv * delta_col; - state[i] = sv.x; state[i+1] = sv.y; state[i+2] = sv.z; state[i+3] = sv.w; + state[i] = FLOAT_TYPE(sv.x); state[i+1] = FLOAT_TYPE(sv.y); + state[i+2] = FLOAT_TYPE(sv.z); state[i+3] = FLOAT_TYPE(sv.w); attn_col += dot(sv, vec4(s_q[i], s_q[i+1], s_q[i+2], s_q[i+3])); } - data_dst[attn_off + col] = attn_col * scale; + data_dst[attn_off + col] = D_TYPE(attn_col * scale); } attn_off += S_V * H; @@ -123,6 +128,6 @@ void main() { } [[unroll]] for (uint i = 0; i < S_V; i++) { - data_dst[s_off + state_base + i * S_V + col] = state[i]; + data_dst[s_off + state_base + i * S_V + col] = D_TYPE(state[i]); } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 4b00ba3debb8..96de0f84da7f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -987,7 +987,8 @@ void process_shaders() { string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); - string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}})); + string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}})); + string_to_spv("gated_delta_net_f16", "gated_delta_net.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float16_t"}, {"FLOAT16", "1"}})); string_to_spv("opt_step_adamw_f32", "opt_step_adamw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); string_to_spv("opt_step_sgd_f32", "opt_step_sgd.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));