vulkan: change gated_delta_net to shard a column across a subgroup (#20662)
* vulkan: change gated_delta_net to shard a column across a subgroup This is based on https://github.com/ggml-org/llama.cpp/pull/20391, I used an LLM to port the CUDA code to Vulkan, and guided to it to make various fixes to work with Vulkan (e.g. handling different subgroup sizes, unknown mapping of subgroup to invocation id, using subgroupAdd optionally, etc.). This fixes a perf regression from the transposing of the values in memory (!20443). * vulkan: Spread columns across fewer lanes to reduce the number of workgroups
This commit is contained in:
@@ -4604,12 +4604,42 @@ 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 bool use_subgroup_reduce = device->subgroup_arithmetic;
|
||||
for (uint32_t si = 0; si < 3; si++) {
|
||||
const uint32_t S_V = gdn_sizes[si];
|
||||
GGML_ASSERT(is_pow2(S_V));
|
||||
|
||||
uint32_t lanes_per_column;
|
||||
if (S_V >= 128u && device->subgroup_clustered) {
|
||||
lanes_per_column = 8u;
|
||||
} else {
|
||||
// Use largest power-of-two that divides both S_V and subgroup_size so that
|
||||
// (1) S_V % lanes_per_column == 0 and (2) S_V % (subgroup_size / lanes_per_column) == 0.
|
||||
// This means we don't need extra bounds checking logic in the shader.
|
||||
lanes_per_column = std::min(S_V, device->subgroup_size);
|
||||
}
|
||||
|
||||
const bool need_clustered_shader = lanes_per_column != 1 && (lanes_per_column < device->subgroup_size);
|
||||
size_t gdn_len;
|
||||
const void * gdn_data;
|
||||
if (use_subgroup_reduce && need_clustered_shader) {
|
||||
gdn_len = gated_delta_net_f32_len;
|
||||
gdn_data = (const void *)gated_delta_net_f32_data;
|
||||
} else if (use_subgroup_reduce) {
|
||||
gdn_len = gated_delta_net_f32_nocluster_len;
|
||||
gdn_data = (const void *)gated_delta_net_f32_nocluster_data;
|
||||
} else {
|
||||
gdn_len = gated_delta_net_f32_shmem_len;
|
||||
gdn_data = (const void *)gated_delta_net_f32_shmem_data;
|
||||
}
|
||||
|
||||
const uint32_t cols_per_wg = device->subgroup_size / lanes_per_column;
|
||||
const std::array<uint32_t, 3> wg_denoms = {1u, 1u, cols_per_wg};
|
||||
|
||||
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);
|
||||
gdn_names[si][kda], gdn_len, gdn_data, "main", 7, sizeof(vk_op_gated_delta_net_push_constants),
|
||||
wg_denoms, {S_V, kda, device->subgroup_size, lanes_per_column}, 1, true, use_subgroup_reduce, device->subgroup_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -10438,7 +10468,7 @@ static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& s
|
||||
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
|
||||
{src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], dst_buf},
|
||||
pc, { H, n_seqs, 1u });
|
||||
pc, { H, n_seqs, S_v });
|
||||
}
|
||||
|
||||
static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
|
||||
|
||||
Reference in New Issue
Block a user