diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 62ef7bb8a2d9..a2d39a37885a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -815,6 +815,11 @@ struct vk_device_struct { bool subgroup_clustered; bool subgroup_vote; bool multi_add; + // DeepSeek-V4 hyper-connection ops, per-op so a kernel can be bisected + // against the unfused graph. + bool dsv4_hc_comb; + bool dsv4_hc_pre; + bool dsv4_hc_post; bool shader_int64; bool buffer_device_address; bool vulkan_memory_model; @@ -1047,6 +1052,9 @@ struct vk_device_struct { vk_pipeline pipeline_cumsum_multipass2_f32; vk_pipeline pipeline_argmax_f32; vk_pipeline pipeline_count_equal_i32; + vk_pipeline pipeline_dsv4_hc_comb_f32; + vk_pipeline pipeline_dsv4_hc_pre_f32; + vk_pipeline pipeline_dsv4_hc_post_f32; std::map pipeline_solve_tri_f32; vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16; @@ -1402,6 +1410,53 @@ struct vk_op_fwht_push_constants { float scale; }; +struct vk_op_dsv4_hc_comb_push_constants { + uint32_t n_tokens; + + uint32_t nbm0; uint32_t nbm1; + uint32_t nbs0; + uint32_t nbb0; + uint32_t nbd0; uint32_t nbd1; uint32_t nbd2; + + uint32_t m_offset; + uint32_t s_offset; + uint32_t b_offset; + uint32_t d_offset; + + float eps; + uint32_t n_iter; +}; + +struct vk_op_dsv4_hc_pre_push_constants { + uint32_t n_embd; + uint32_t n_tokens; + + uint32_t nbx0; uint32_t nbx1; uint32_t nbx2; + uint32_t nbw0; uint32_t nbw1; + uint32_t nbd0; uint32_t nbd1; + + uint32_t x_offset; + uint32_t w_offset; + uint32_t d_offset; +}; + +struct vk_op_dsv4_hc_post_push_constants { + uint32_t n_embd; + uint32_t n_tokens; + + uint32_t nbx0; uint32_t nbx1; + uint32_t nbr0; uint32_t nbr1; uint32_t nbr2; + uint32_t nbp0; uint32_t nbp1; + uint32_t nbc0; uint32_t nbc1; uint32_t nbc2; + uint32_t nbd0; uint32_t nbd1; uint32_t nbd2; + + uint32_t x_offset; + uint32_t r_offset; + uint32_t p_offset; + uint32_t c_offset; + uint32_t d_offset; +}; + struct vk_op_count_experts_push_constants { uint32_t ne00; uint32_t ne01; @@ -2498,6 +2553,32 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk GGML_UNUSED(src3); } +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_comb_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { + p.m_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + p.s_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + p.b_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type); + p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + GGML_UNUSED(src3); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_pre_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { + p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + p.w_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + GGML_UNUSED(src2); + GGML_UNUSED(src3); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_post_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { + p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + p.r_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + p.p_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type); + p.c_offset = get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type); + p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); +} + struct ggml_backend_vk_buffer_context { vk_device_ref device; vk_buffer dev_buffer; @@ -5784,6 +5865,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); + // comb holds a token's 4x4 matrix in one 16-lane slice of a subgroup, so it + // needs at least 16 lanes, pinned to a known size. + if (device->subgroup_basic && device->subgroup_shuffle && device->subgroup_require_full_support && device->subgroup_size >= 16) { + const uint32_t tokens_per_workgroup = 4 * (device->subgroup_size / 16); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, sizeof(vk_op_dsv4_hc_comb_push_constants), {tokens_per_workgroup, 1, 1}, { device->subgroup_size }, 1, true, true, device->subgroup_size); + } + + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32, "dsv4_hc_pre_f32", dsv4_hc_pre_f32_len, dsv4_hc_pre_f32_data, "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants), {256, 1, 1}, { 256 }, 1); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256 }, 1); + for (auto &s : device->pipeline_solve_tri_f32) { const vk_solve_tri_pipeline_state &state = s.first; @@ -6679,6 +6770,11 @@ static vk_device ggml_vk_get_device(size_t idx) { device->properties.limits.maxPushConstantsSize >= sizeof(vk_op_multi_add_push_constants) && getenv("GGML_VK_DISABLE_MULTI_ADD") == nullptr; + const bool dsv4_hc_all = getenv("GGML_VK_DISABLE_DSV4_HC") == nullptr; + device->dsv4_hc_comb = dsv4_hc_all && getenv("GGML_VK_DISABLE_DSV4_HC_COMB") == nullptr; + device->dsv4_hc_pre = dsv4_hc_all && getenv("GGML_VK_DISABLE_DSV4_HC_PRE") == nullptr; + device->dsv4_hc_post = dsv4_hc_all && getenv("GGML_VK_DISABLE_DSV4_HC_POST") == nullptr; + device->shader_int64 = device_features2.features.shaderInt64; device->buffer_device_address = vk12_features.bufferDeviceAddress; device->vulkan_memory_model = vk12_features.vulkanMemoryModel; @@ -9999,6 +10095,122 @@ static void ggml_vk_fwht(ggml_backend_vk_context * ctx, vk_context& subctx, cons ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, dst_buf }, pc, { workgroups_x, 1, 1 }); } +// Element stride; supports_op checks that the byte strides divide evenly. +static uint32_t ggml_vk_nb_elem(const ggml_tensor * t, int i) { + return (uint32_t)(t->nb[i] / ggml_type_size(t->type)); +} + +// The HC shaders address their tensors in whole f32 elements. +static bool ggml_vk_dsv4_hc_strides_ok(const ggml_tensor * op) { + auto const & strides_ok = [](const ggml_tensor * t) { + const size_t ts = ggml_type_size(t->type); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (t->nb[i] % ts != 0) { + return false; + } + } + return true; + }; + + if (!strides_ok(op)) { + return false; + } + for (uint32_t i = 0; i < GGML_MAX_SRC; ++i) { + if (op->src[i] && !strides_ok(op->src[i])) { + return false; + } + } + return true; +} + +static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * mixes, const ggml_tensor * scale, const ggml_tensor * base, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_vk_dsv4_hc_comb(" << mixes << ", " << scale << ", " << base << ", " << dst << ")"); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_comb_f32; + GGML_ASSERT(pipeline != nullptr); + + const uint32_t n_tokens = (uint32_t)mixes->ne[1]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_subbuffer mixes_buf = ggml_vk_tensor_subbuffer(ctx, mixes, true); + const vk_subbuffer scale_buf = ggml_vk_tensor_subbuffer(ctx, scale, true); + const vk_subbuffer base_buf = ggml_vk_tensor_subbuffer(ctx, base, true); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_op_dsv4_hc_comb_push_constants pc = { + n_tokens, + ggml_vk_nb_elem(mixes, 0), ggml_vk_nb_elem(mixes, 1), + ggml_vk_nb_elem(scale, 0), + ggml_vk_nb_elem(base, 0), + ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), + 0, 0, 0, 0, + ggml_get_op_params_f32(dst, 0), + (uint32_t)ggml_get_op_params_i32(dst, 1), + }; + init_pushconst_tensor_offsets(ctx, pc, mixes, scale, base, nullptr, dst); + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { mixes_buf, scale_buf, base_buf, dst_buf }, pc, { n_tokens, 1, 1 }); +} + +static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_vk_dsv4_hc_pre(" << x << ", " << weights << ", " << dst << ")"); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32; + GGML_ASSERT(pipeline != nullptr); + + const uint32_t n_embd = (uint32_t)x->ne[0]; + const uint32_t n_tokens = (uint32_t)x->ne[2]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true); + const vk_subbuffer w_buf = ggml_vk_tensor_subbuffer(ctx, weights, true); + const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_op_dsv4_hc_pre_push_constants pc = { + n_embd, n_tokens, + ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), ggml_vk_nb_elem(x, 2), + ggml_vk_nb_elem(weights, 0), ggml_vk_nb_elem(weights, 1), + ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), + 0, 0, 0, + }; + init_pushconst_tensor_offsets(ctx, pc, x, weights, nullptr, nullptr, dst); + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, w_buf, d_buf }, pc, { n_embd, n_tokens, 1 }); +} + +static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_vk_dsv4_hc_post(" << x << ", " << residual << ", " << post << ", " << comb << ", " << dst << ")"); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32; + GGML_ASSERT(pipeline != nullptr); + + const uint32_t n_embd = (uint32_t)x->ne[0]; + const uint32_t n_tokens = (uint32_t)x->ne[1]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true); + const vk_subbuffer r_buf = ggml_vk_tensor_subbuffer(ctx, residual, true); + const vk_subbuffer p_buf = ggml_vk_tensor_subbuffer(ctx, post, true); + const vk_subbuffer c_buf = ggml_vk_tensor_subbuffer(ctx, comb, true); + const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_op_dsv4_hc_post_push_constants pc = { + n_embd, n_tokens, + ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), + ggml_vk_nb_elem(residual, 0), ggml_vk_nb_elem(residual, 1), ggml_vk_nb_elem(residual, 2), + ggml_vk_nb_elem(post, 0), ggml_vk_nb_elem(post, 1), + ggml_vk_nb_elem(comb, 0), ggml_vk_nb_elem(comb, 1), ggml_vk_nb_elem(comb, 2), + ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), + 0, 0, 0, 0, 0, + }; + init_pushconst_tensor_offsets(ctx, pc, x, residual, post, comb, dst); + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, r_buf, p_buf, c_buf, d_buf }, pc, { n_embd, n_tokens, 1 }); +} + static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { ggml_tensor * dst = cgraph->nodes[node_idx]; ggml_tensor * src0 = dst->src[0]; @@ -15588,6 +15800,18 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_OP_CUMSUM: ggml_vk_cumsum(ctx, compute_ctx, src0, node); + break; + case GGML_OP_DSV4_HC_COMB: + ggml_vk_dsv4_hc_comb(ctx, compute_ctx, src0, src1, src2, node); + + break; + case GGML_OP_DSV4_HC_PRE: + ggml_vk_dsv4_hc_pre(ctx, compute_ctx, src0, src1, node); + + break; + case GGML_OP_DSV4_HC_POST: + ggml_vk_dsv4_hc_post(ctx, compute_ctx, src0, src1, src2, src3, node); + break; case GGML_OP_MEAN: ggml_vk_mean(ctx, compute_ctx, src0, node); @@ -18399,6 +18623,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } return false; } + case GGML_OP_DSV4_HC_COMB: + case GGML_OP_DSV4_HC_PRE: + case GGML_OP_DSV4_HC_POST: + { + const bool enabled = + op->op == GGML_OP_DSV4_HC_COMB ? device->dsv4_hc_comb : + op->op == GGML_OP_DSV4_HC_PRE ? device->dsv4_hc_pre : + device->dsv4_hc_post; + if (!enabled || op->type != GGML_TYPE_F32) { + return false; + } + for (uint32_t i = 0; i < GGML_MAX_SRC; ++i) { + if (op->src[i] && op->src[i]->type != GGML_TYPE_F32) { + return false; + } + } + if (!ggml_vk_dsv4_hc_strides_ok(op)) { + return false; + } + // hc is hardcoded to 4 in the shaders. ggml only constrains it + // to 4 for COMB, so PRE/POST have to be checked here. + if (op->op == GGML_OP_DSV4_HC_PRE && op->src[0]->ne[1] != 4) { + return false; + } + if (op->op == GGML_OP_DSV4_HC_POST && op->src[1]->ne[1] != 4) { + return false; + } + if (op->op == GGML_OP_DSV4_HC_COMB) { + return device->pipeline_dsv4_hc_comb_f32 != nullptr; + } + // pre/post launch one workgroup row per token + const uint32_t n_tokens = (uint32_t)(op->op == GGML_OP_DSV4_HC_PRE ? op->src[0]->ne[2] : op->src[0]->ne[1]); + return n_tokens <= device->properties.limits.maxComputeWorkGroupCount[1]; + } case GGML_OP_SOLVE_TRI: { if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32) { @@ -19339,6 +19597,13 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * tensor_clone = ggml_sum_rows(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_CUMSUM) { tensor_clone = ggml_cumsum(ggml_ctx, src_clone[0]); + } else if (tensor->op == GGML_OP_DSV4_HC_COMB) { + tensor_clone = ggml_dsv4_hc_comb(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], + ggml_get_op_params_f32(tensor, 0), ggml_get_op_params_i32(tensor, 1)); + } else if (tensor->op == GGML_OP_DSV4_HC_PRE) { + tensor_clone = ggml_dsv4_hc_pre(ggml_ctx, src_clone[0], src_clone[1]); + } else if (tensor->op == GGML_OP_DSV4_HC_POST) { + tensor_clone = ggml_dsv4_hc_post(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]); } else if (tensor->op == GGML_OP_MEAN) { tensor_clone = ggml_mean(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_ARGMAX) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp new file mode 100644 index 000000000000..5bc24f5859de --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp @@ -0,0 +1,94 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require + +// Each token's 4x4 comb matrix lives in 16 consecutive subgroup lanes, indexed +// idst + hc*isrc to match the CPU reference. That puts idst in bits 0..1 and +// isrc in bits 2..3, so subgroupShuffleXor by 1|2 reduces a row and by 4|8 +// reduces a column. + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 32; + +layout(local_size_x_id = 0, local_size_y = 4, local_size_z = 1) in; + +layout(push_constant) uniform parameter +{ + uint n_tokens; + + uint nbm0; uint nbm1; // mixes + uint nbs0; // scale + uint nbb0; // base + uint nbd0; uint nbd1; uint nbd2; // dst + + uint m_offset; + uint s_offset; + uint b_offset; + uint d_offset; + + float eps; + uint n_iter; +}; + +layout(binding = 0, std430) readonly buffer M { float data_m[]; }; +layout(binding = 1, std430) readonly buffer S { float data_s[]; }; +layout(binding = 2, std430) readonly buffer B { float data_b[]; }; +layout(binding = 3, std430) writeonly buffer D { float data_d[]; }; + +const uint hc = 4; +const uint comb_offset = 2 * hc; + +const uint TOKENS_PER_SUBGROUP = SUBGROUP_SIZE / 16; + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint blk = lane >> 4; // which 16-lane block, i.e. which token + const uint idx = lane & 15; // idst + hc*isrc + + const uint sg = gl_WorkGroupID.x * gl_WorkGroupSize.y + gl_SubgroupID; + const uint it = sg * TOKENS_PER_SUBGROUP + blk; + + // No early return: every lane must stay active through the last shuffle. + // Out-of-range blocks compute a discarded value; the masks stay inside a + // 16-lane block, so they cannot contaminate a live token. + const bool in_range = it < n_tokens; + + const float scale_comb = data_s[s_offset + 2 * nbs0]; + + float v = 0.0f; + if (in_range) { + v = data_m[m_offset + (comb_offset + idx) * nbm0 + it * nbm1] * scale_comb + + data_b[b_offset + (comb_offset + idx) * nbb0]; + } + + // Softmax across destinations: the four lanes sharing an isrc. + float vmax = max(v, subgroupShuffleXor(v, 1)); + vmax = max(vmax, subgroupShuffleXor(vmax, 2)); + v = exp(v - vmax); + + float sum = v + subgroupShuffleXor(v, 1); + sum += subgroupShuffleXor(sum, 2); + v = v / sum + eps; + + // Normalize columns: equal destination indices are four lanes apart. + sum = v + subgroupShuffleXor(v, 4); + sum += subgroupShuffleXor(sum, 8); + v /= sum + eps; + + for (uint i = 1; i < n_iter; ++i) { + sum = v + subgroupShuffleXor(v, 1); + sum += subgroupShuffleXor(sum, 2); + v /= sum + eps; + + sum = v + subgroupShuffleXor(v, 4); + sum += subgroupShuffleXor(sum, 8); + v /= sum + eps; + } + + if (in_range) { + const uint idst = idx & 3; + const uint isrc = idx >> 2; + data_d[d_offset + idst * nbd0 + isrc * nbd1 + it * nbd2] = v; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp new file mode 100644 index 000000000000..bab6f8767848 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp @@ -0,0 +1,83 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +// Fan one stream back out to hc streams and add the combination-weighted +// residuals: +// +// dst[i0, idst, it] = x[i0, it]*post[idst, it] +// + sum_isrc residual[i0, isrc, it]*comb[idst, isrc, it] + +layout(constant_id = 0) const uint BLOCK_SIZE = 256; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform parameter +{ + uint n_embd; + uint n_tokens; + + uint nbx0; uint nbx1; // x + uint nbr0; uint nbr1; uint nbr2; // residual + uint nbp0; uint nbp1; // post + uint nbc0; uint nbc1; uint nbc2; // comb + uint nbd0; uint nbd1; uint nbd2; // dst + + uint x_offset; + uint r_offset; + uint p_offset; + uint c_offset; + uint d_offset; +}; + +layout(binding = 0, std430) readonly buffer X { float data_x[]; }; +layout(binding = 1, std430) readonly buffer R { float data_r[]; }; +layout(binding = 2, std430) readonly buffer P { float data_p[]; }; +layout(binding = 3, std430) readonly buffer C { float data_c[]; }; +layout(binding = 4, std430) writeonly buffer D { float data_d[]; }; + +const uint hc = 4; + +shared float post_s[hc]; +shared float comb_s[hc * hc]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint it = gl_WorkGroupID.y; + + if (tid < hc) { + post_s[tid] = data_p[p_offset + tid * nbp0 + it * nbp1]; + } + if (tid < hc * hc) { + const uint idst = tid & 3; + const uint isrc = tid >> 2; + comb_s[tid] = data_c[c_offset + idst * nbc0 + isrc * nbc1 + it * nbc2]; + } + barrier(); + + // After the barrier, so every invocation reaches it. + const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid; + if (i0 >= n_embd) { + return; + } + + const float xv = data_x[x_offset + i0 * nbx0 + it * nbx1]; + + const uint rb = r_offset + i0 * nbr0 + it * nbr2; + + float r[hc]; + [[unroll]] + for (uint isrc = 0; isrc < hc; ++isrc) { + r[isrc] = data_r[rb + isrc * nbr1]; + } + + [[unroll]] + for (uint idst = 0; idst < hc; ++idst) { + float result = xv * post_s[idst]; + [[unroll]] + for (uint isrc = 0; isrc < hc; ++isrc) { + result = fma(r[isrc], comb_s[idst + hc * isrc], result); + } + data_d[d_offset + i0 * nbd0 + idst * nbd1 + it * nbd2] = result; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp new file mode 100644 index 000000000000..51deabbac6ed --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp @@ -0,0 +1,59 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +// Collapse the hc residual streams of a token into one, weighted per stream: +// +// dst[i0, it] = sum_ih x[i0, ih, it] * weights[ih, it] + +layout(constant_id = 0) const uint BLOCK_SIZE = 256; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform parameter +{ + uint n_embd; + uint n_tokens; + + uint nbx0; uint nbx1; uint nbx2; // x + uint nbw0; uint nbw1; // weights + uint nbd0; uint nbd1; // dst + + uint x_offset; + uint w_offset; + uint d_offset; +}; + +layout(binding = 0, std430) readonly buffer X { float data_x[]; }; +layout(binding = 1, std430) readonly buffer W { float data_w[]; }; +layout(binding = 2, std430) writeonly buffer D { float data_d[]; }; + +const uint hc = 4; + +shared float w[hc]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint it = gl_WorkGroupID.y; + + if (tid < hc) { + w[tid] = data_w[w_offset + tid * nbw0 + it * nbw1]; + } + barrier(); + + // After the barrier, so every invocation reaches it. + const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid; + if (i0 >= n_embd) { + return; + } + + const uint xb = x_offset + i0 * nbx0 + it * nbx2; + + float result = 0.0f; + [[unroll]] + for (uint ih = 0; ih < hc; ++ih) { + result = fma(data_x[xb + ih * nbx1], w[ih], result); + } + + data_d[d_offset + i0 * nbd0 + it * nbd1] = result; +} 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 fbc2ea3ca226..8b7d14c5272a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1027,6 +1027,9 @@ void process_shaders() { string_to_spv("fwht_f32", "fwht.comp", {}); string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}}); string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}})); + string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); + string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); + string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); string_to_spv("cumsum_f32", "cumsum.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("cumsum_multipass1_f32", "cumsum_multipass1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 2a1851b59d5e..3ba4a4b4e341 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8351,6 +8351,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 20)); + // production n_iter (DeepSeek-V4 uses 20) across batch sizes that cross + // subgroup and workgroup boundaries; 1 = single-token decode + for (int64_t n_tokens : {1, 256, 336, 512, 513, 1024, 2048}) { + test_cases.emplace_back(new test_dsv4_hc_comb(n_tokens, 20)); + } test_cases.emplace_back(new test_dsv4_hc_pre(1, 1)); test_cases.emplace_back(new test_dsv4_hc_pre(31, 17));