Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-04 02:37:27 +02:00
Code Issues 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 <[email protected]>
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) {
ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp
+93 -2
View File
@@ -2,6 +2,11 @@
#extension GL_EXT_control_flow_attributes : enable
#ifdef USE_SUBGROUPS
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#endif
#include "types.glsl"
layout (push_constant) uniform parameter
@@ -11,6 +16,10 @@ layout (push_constant) uniform parameter
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;
} p;
#define BLOCK_SIZE 256
@@ -21,16 +30,98 @@ layout (binding = 0) readonly buffer A {uint data_a[];};
layout (binding = 1) writeonly buffer D {uint data_d[];};
shared uint vals[BLOCK_SIZE];
shared uint offsets[BLOCK_SIZE];
shared uint cursors[BLOCK_SIZE];
// see init_fastdiv_values in ggml-vulkan.cpp
uint fastdiv(uint n, uint mp, uint L) {
uint msbs, lsbs;
// msbs = mulhi(n, mp)
umulExtended(n, mp, msbs, lsbs);
return (msbs + n) >> L;
}
// data_d layout when p.hoist_row_ids is set:
// [0, n_experts) per-expert row count
// [n_experts, 2*n_experts) per-expert start offset into the row id region
// [2*n_experts] total row count
// [2*n_experts + 1, ) row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff)
// Otherwise only data_d[expert_id] is written, holding that expert's row count.
void main() {
const uint expert_id = gl_WorkGroupID.x;
const uint num_elements = p.ne00 * p.ne01;
const uint tid = gl_LocalInvocationID.x;
if (p.hoist_row_ids != 0) {
if (tid < p.n_experts) {
vals[tid] = 0;
}
barrier();
for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
const uint i00 = idx - i01 * p.ne00;
const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
if (expert < p.n_experts) {
atomicAdd(vals[expert], 1);
}
}
barrier();
#ifdef USE_SUBGROUPS
if (gl_SubgroupID == 0) {
// pad the trip count so the subgroup ops stay in uniform control flow
const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1);
uint base = 0;
for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) {
const bool in_range = expert < p.n_experts;
const uint count = in_range ? vals[expert] : 0;
const uint offset = base + subgroupExclusiveAdd(count);
if (in_range) {
data_d[expert] = count;
data_d[p.n_experts + expert] = offset;
offsets[expert] = offset;
cursors[expert] = 0;
}
base += subgroupAdd(count);
}
if (subgroupElect()) {
data_d[2 * p.n_experts] = base;
}
}
#else
if (tid == 0) {
uint offset = 0;
for (uint expert = 0; expert < p.n_experts; ++expert) {
const uint count = vals[expert];
data_d[expert] = count;
data_d[p.n_experts + expert] = offset;
offsets[expert] = offset;
cursors[expert] = 0;
offset += count;
}
data_d[2 * p.n_experts] = offset;
}
#endif
barrier();
for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
const uint i00 = idx - i01 * p.ne00;
const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
if (expert < p.n_experts) {
const uint row = atomicAdd(cursors[expert], 1);
const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu);
data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id;
}
}
return;
}
uint count = 0;
for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
const uint i01 = idx / p.ne00;
const uint i00 = idx % p.ne00;
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
const uint i00 = idx - i01 * p.ne00;
const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
count += uint(a == expert_id);
ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp
+21 -14
View File
@@ -88,6 +88,9 @@ layout (push_constant) uniform parameter
uint nei1;
uint nbi1;
uint ne11;
uint padded_N;
uint n_experts;
uint hoist_row_ids;
#else
uint base_work_group_z;
uint num_batches;
@@ -214,27 +217,31 @@ void main() {
const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK;
#ifdef MUL_MAT_ID
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
if (p.hoist_row_ids != 0) {
load_row_ids_hoisted(expert_idx, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#else
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
}
_ne1++;
}
_ne1++;
}
}
}
barrier();
barrier();
#endif
}
// Workgroup has no work
if (ic * BN >= _ne1) return;
ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
+24 -1
View File
@@ -67,6 +67,10 @@ layout (push_constant) uniform parameter
#endif
// N dimension for the B matrix can be >= p.N
uint padded_N;
#ifdef MUL_MAT_ID
uint n_experts;
uint hoist_row_ids;
#endif
} p;
@@ -225,6 +229,23 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
}
barrier();
}
void load_row_ids_hoisted(uint expert_idx, uint ic) {
_ne1 = uint(data_expert_count[expert_idx]);
const uint tile_begin = ic * BN;
const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
const uint ii0 = packed_row_id & 0xffffu;
const uint ii1 = packed_row_id >> 16;
row_ids[i] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0);
}
barrier();
}
#endif
void main() {
@@ -266,7 +287,9 @@ void main() {
const uint ik = gl_WorkGroupID.x / blocks_m;
#ifdef MUL_MAT_ID
if (bitCount(p.nei0) == 1) {
if (p.hoist_row_ids != 0) {
load_row_ids_hoisted(expert_idx, ic);
} else if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
} else {
load_row_ids(expert_idx, false, ic);
ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl
+15
View File
@@ -71,4 +71,19 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
barrier();
}
#endif // MUL_MAT_ID_USE_SUBGROUPS
void load_row_ids_hoisted(uint expert_idx, uint ic) {
_ne1 = uint(data_expert_count[expert_idx]);
const uint tile_begin = ic * BN;
const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
row_ids[i] = u16vec2(packed_row_id & 0xffffu, packed_row_id >> 16);
}
barrier();
}
#endif // MUL_MAT_ID
ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp
+21 -14
View File
@@ -56,6 +56,9 @@ layout (push_constant) uniform parameter
uint nei1;
uint nbi1;
uint ne11;
uint padded_N;
uint n_experts;
uint hoist_row_ids;
#else
uint base_work_group_z;
uint num_batches;
@@ -157,27 +160,31 @@ void main() {
const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK;
#ifdef MUL_MAT_ID
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
if (p.hoist_row_ids != 0) {
load_row_ids_hoisted(expert_idx, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#else
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
}
_ne1++;
}
_ne1++;
}
}
}
barrier();
barrier();
#endif
}
// Workgroup has no work
if (ic * BN >= _ne1) return;
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+1
View File
@@ -1039,6 +1039,7 @@ void process_shaders() {
string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("count_experts", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}));
string_to_spv("count_experts_subgroup", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}, {"USE_SUBGROUPS", "1"}}));
for (std::string dim_str : {"", "_3d"}) {
for (bool bda : {false, true}) {