mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-04 02:37:27 +02:00
CUDA: use the fast mm_ids_helper path for any n_expert_used (#27978)
The optimized path grouped warp lanes by token and required warp_size % n_expert_used == 0, with a single hardcoded exception padding 6 up to 8. Every other count fell back to the generic path, which walks the tokens one at a time with a warp reduction per token, for each of the n_expert blocks. The lane group only has to divide the warp, and the loop body already guards the padded lanes with iex < n_expert_used, so the padding generalizes to the next power of two. The 6 -> 8 case and every count already dispatched keep the exact same padding as before. n_expert_used = 10 now reaches the fast path. Measured on Qwen3.8-Flash-Next (512 experts, 10 used) at 55k context on an RTX PRO 6000, warm runs with the first one discarded: prompt processing 2334 -> 2600 t/s Token generation is unaffected, since a single token leaves nothing to walk. Other expert counts reach the fast path by adding their case to the dispatch.
This commit is contained in:
@@ -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 <int n> 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<n_expert_used_template>::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;
|
||||
|
||||
Reference in New Issue
Block a user