diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 72e844aebf..28b2d875d2 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -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 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) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp index ffc8608691..83c56fce52 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp @@ -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); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 3df88044a5..c1ccac7aa2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -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; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index a2e15f6f5c..cf78474a92 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -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); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl index 26c5c12a49..54ad60b2ef 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl @@ -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 diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp index aae1c2e8ae..c2d84c05b4 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp @@ -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; 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 0da943da95..d375c2d127 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -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}) {