Vulkan: add hoisting support for row IDs and expert count in shaders

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -1327,6 +1327,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;
@@ -1407,6 +1409,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 {
@@ -1585,6 +1591,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;
@@ -5754,7 +5764,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;
@@ -8869,13 +8883,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 });
 }
 
@@ -10049,6 +10063,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;
@@ -10189,7 +10209,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 (
@@ -10205,8 +10226,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);
         }
 
@@ -10272,18 +10293,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) {
@@ -10352,7 +10378,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

@@ -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

@@ -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;
@@ -211,27 +214,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

@@ -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

@@ -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