diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu index f80442fbe4..ed0851dcf8 100644 --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu @@ -19,6 +19,11 @@ struct mm_ids_helper_store { }; static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store"); +// the generic path passes 0, which needs no padding since it never groups lanes by token +template struct mm_ids_pow2 { static constexpr int value = 2*mm_ids_pow2<(n + 1)/2>::value; }; +template <> struct mm_ids_pow2<1> { static constexpr int value = 1; }; +template <> struct mm_ids_pow2<0> { static constexpr int value = 1; }; + // Helper function for mul_mat_id, converts ids to a more convenient format. // ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert. // ids_dst describes the same mapping but for the dst tensor. @@ -32,6 +37,9 @@ static __global__ void mm_ids_helper( const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template; const int expert = blockIdx.x; + // token slots per warp lane group, padded to a power of 2 so a warp divides evenly + constexpr int neu_padded = mm_ids_pow2::value; + extern __shared__ char data_mm_ids_helper[]; mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper; @@ -60,8 +68,8 @@ static __global__ void mm_ids_helper( } } else { // Implementation optimized for specific numbers of experts used: - static_assert(n_expert_used == 6 || warp_size % n_expert_used == 0, "bad n_expert_used"); - const int neu_padded = n_expert_used == 6 ? 8 : n_expert_used; // Padded to next higher power of 2. + // a warp holds a whole number of token slots, so the slot count is padded to a power of 2 + static_assert(neu_padded <= warp_size && warp_size % neu_padded == 0, "bad n_expert_used"); for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) { const int it = it0 + threadIdx.x / neu_padded; @@ -156,6 +164,9 @@ void ggml_cuda_launch_mm_ids_helper( case 8: launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); break; + case 10: + launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); + break; case 16: launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); break;