Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Vulkan: add hoisting support for row IDs and expert count in shaders (#26686)

Browse Source 
* vulkan: add hoisting support for row IDs and expert count in shaders

* use hoisted row ids in coopmat2

* vulkan: address review feedback on count_experts
- use vk_op_count_experts_push_constants instead of a raw uint vector
- apply the fastdiv trick to the ne00 div/mod in count_experts
- compute the per-expert offsets with subgroupExclusiveAdd when the
  device supports it, keeping the serial path as fallback
- document the data_d layout and the hoisted_row_id_words bound
- drop a leftover debug print in ggml_vk_matmul_id

* vulkan: use init_pushconst_fastdiv for count_experts push constants

* vulkan: refine comments for row ID hoisting and data layout in count_experts shader

* Whitespace

---------

Co-authored-by: Jeff Bolz <jbolz@nvidia.com>
This commit is contained in:
ravel7524
2026-08-28 16:52:49 +02:00
committed by GitHub
co-authored by Jeff Bolz
parent 8663224818
commit 90c26fcd4b
7 changed files with 212 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+37 -11
View File
@@ -1348,6 +1348,8 @@ struct vk_mat_mat_id_push_constants {
uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d;
uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11;
uint32_t padded_N;
uint32_t n_experts;
uint32_t hoist_row_ids;
};
struct vk_mat_vec_id_push_constants {
uint32_t ncols;
@@ -1428,6 +1430,10 @@ struct vk_op_count_experts_push_constants {
uint32_t nb00;
uint32_t nb01;
uint32_t a_offset;
uint32_t n_experts;
uint32_t hoist_row_ids;
uint32_t ne00mp;
uint32_t ne00L;
};
struct vk_op_glu_push_constants {
@@ -1606,6 +1612,10 @@ template <> void init_pushconst_fastdiv(vk_op_glu_push_constants &p) {
init_fastdiv_values(p.ne20, p.ne2_0mp, p.ne2_0L);
}
template <> void init_pushconst_fastdiv(vk_op_count_experts_push_constants &p) {
init_fastdiv_values(p.ne00, p.ne00mp, p.ne00L);
}
struct vk_op_binary_push_constants {
uint32_t ne;
uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03;
@@ -5839,7 +5849,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1);
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);
if (device->subgroup_arithmetic && device->subgroup_require_full_support) {
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
} else {
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);
}
for (auto &s : device->pipeline_solve_tri_f32) {
const vk_solve_tri_pipeline_state &state = s.first;
@@ -8970,13 +8984,13 @@ static void ggml_vk_matmul_id(
uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d,
uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d,
uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11,
uint32_t padded_n) {
uint32_t padded_n, bool hoist_row_ids) {
VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " <<
"m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " <<
"batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " <<
"n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")");
const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d,
nei0, nei1, nbi1, ne11, padded_n };
nei0, nei1, nbi1, ne11, padded_n, n_as, uint32_t(hoist_row_ids) };
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as });
}
@@ -10162,6 +10176,12 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
// const uint64_t ne23 = dst->ne[3];
const uint64_t n_as = ne02;
// n_as counts, n_as offsets, one total, then one packed row id per (expert, token).
// Hoisting requires 16-bit indices for the packing and a table that fits one binding.
const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1;
const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff &&
hoisted_row_id_words * sizeof(uint32_t) <=
ctx->device->properties.limits.maxStorageBufferRange;
ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context;
@@ -10302,7 +10322,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
}
vk_pipeline count_experts = ctx->device->pipeline_count_experts;
uint32_t expert_count_size = sizeof(uint32_t) * n_as;
const size_t expert_data_size = sizeof(uint32_t) *
(hoist_row_ids ? hoisted_row_id_words : n_as);
{
if (
@@ -10318,8 +10339,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
ctx->prealloc_size_y = y_sz;
ggml_vk_preallocate_buffers(ctx, subctx);
}
if (ctx->prealloc_size_split_k < expert_count_size) {
ctx->prealloc_size_split_k = expert_count_size;
if (ctx->prealloc_size_split_k < expert_data_size) {
ctx->prealloc_size_split_k = expert_data_size;
ggml_vk_preallocate_buffers(ctx, subctx);
}
@@ -10385,18 +10406,23 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
}
}
// Count how many times each expert is used
vk_subbuffer expert_count_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
vk_subbuffer expert_count_buf = { ctx->prealloc_split_k, 0, expert_data_size };
if (ctx->prealloc_split_k_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}
{
const std::vector<uint32_t> pc = { (uint32_t)nei0,
vk_op_count_experts_push_constants pc = { (uint32_t)nei0,
(uint32_t)nei1,
(uint32_t)(nbi0 / ggml_type_size(ids->type)),
(uint32_t)(nbi1 / ggml_type_size(ids->type)),
(uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)) };
(uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)),
(uint32_t)n_as,
uint32_t(hoist_row_ids),
0, 0 };
init_pushconst_fastdiv(pc);
ggml_vk_dispatch_pipeline(ctx, subctx, count_experts,
{ vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, { (uint32_t)n_as, 1, 1});
{ vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc,
{ hoist_row_ids ? 1u : (uint32_t)n_as, 1, 1});
}
if (x_non_contig) {
@@ -10465,7 +10491,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
{ d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf,
ne01, ne21, ne10, ne10, stride_b_y, ne01,
stride_batch_x, stride_batch_y, ne20*ne21,
n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n
n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n, hoist_row_ids
); // NOLINT
if (x_non_contig || qx_needs_dequant) {