Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-20 01:31:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix it

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-06 13:30:01 +02:00
parent f8fa4df82a
commit 1a0e158115
1 changed files with 51 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/mmid.cu
+51 -17
View File
@@ -23,14 +23,21 @@ static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_help
// 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.
// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1].
//
// Normally each token's n_expert_used slots reference distinct experts, so at most one slot per token
// ever lands in a given expert's bucket. A malformed routing table (e.g. from a third-party expert-pruning
// tool, see issue #26588) can select the same expert more than once for one token; the store buffer below
// is sized with a small amount of slack (STORE_CAPACITY_MUL) to keep handling that case memory-safe.
template <int n_expert_used_template>
__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int STORE_CAPACITY_MUL = 2; // must match launch_mm_ids_helper's shared memory allocation
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
const int store_capacity = STORE_CAPACITY_MUL*n_tokens;
extern __shared__ char data_mm_ids_helper[];
mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
@@ -39,23 +46,33 @@ static __global__ void mm_ids_helper(
int it_compact = 0; // Running index for the compact slice of this expert.
if constexpr (n_expert_used_template == 0) {
// Generic implementation:
// Generic implementation, one warp-synchronous step per warp_size experts used:
for (int it = 0; it < n_tokens; ++it) {
int iex_used = -1; // The index at which the expert is used, if any.
for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
const int expert_used = ids[it*si1 + iex];
for (int iex0 = 0; iex0 < n_expert_used; iex0 += warp_size) {
const int iex = iex0 + threadIdx.x;
const int expert_used = iex < n_expert_used ? ids[it*si1 + iex] : INT_MAX;
nex_prev += expert_used < expert;
if (expert_used == expert) {
iex_used = iex;
const int iex_used = expert_used == expert ? iex : -1;
const int is_match = iex_used != -1 ? 1 : 0;
// A token can select the same expert more than once, give each match its own row via
// an inclusive prefix sum of matches over the warp (usually at most 1 match total).
int match_prefix = is_match;
#pragma unroll
for (int offset = 1; offset < warp_size; offset *= 2) {
const int n = __shfl_up_sync(0xFFFFFFFF, match_prefix, offset, warp_size);
if (threadIdx.x >= static_cast<unsigned int>(offset)) {
match_prefix += n;
}
}
}
if (iex_used != -1) {
store[it_compact] = mm_ids_helper_store(it, iex_used);
}
const int idx = it_compact + match_prefix - 1;
if (iex_used != -1 && idx < store_capacity) {
store[idx] = mm_ids_helper_store(it, iex_used);
}
if (warp_reduce_any<warp_size>(iex_used != -1)) {
it_compact++;
it_compact += __shfl_sync(0xFFFFFFFF, match_prefix, warp_size - 1, warp_size);
}
}
} else {
@@ -71,8 +88,18 @@ static __global__ void mm_ids_helper(
const int iex_used = expert_used == expert ? iex : -1;
nex_prev += expert_used < expert;
// Whether the threads at this token position have used the expert:
const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);
// A token can select the same expert more than once: count how many of its slots match
// (usually 0 or 1) and give each one its own row via a prefix sum within the token's lane group.
const int is_match = iex_used != -1 ? 1 : 0;
int match_prefix = is_match;
#pragma unroll
for (int offset = 1; offset < neu_padded; offset *= 2) {
const int n = __shfl_up_sync(0xFFFFFFFF, match_prefix, offset, neu_padded);
if (iex >= offset) {
match_prefix += n;
}
}
const int it_compact_add_self = warp_reduce_sum<neu_padded>(is_match); // number of matches for this token.
// Do a scan over threads at lower token positions in warp to get the correct index for writing data:
int it_compact_add_lower = 0;
@@ -84,8 +111,9 @@ static __global__ void mm_ids_helper(
}
}
if (iex_used != -1) {
store[it_compact + it_compact_add_lower] = mm_ids_helper_store(it, iex_used);
const int idx = it_compact + it_compact_add_lower + match_prefix - 1;
if (iex_used != -1 && idx < store_capacity) {
store[idx] = mm_ids_helper_store(it, iex_used);
}
// The thread with the highest index in the warp always has the sum over the whole warp, use it to increment all threads:
@@ -94,6 +122,10 @@ static __global__ void mm_ids_helper(
}
nex_prev = warp_reduce_sum<warp_size>(nex_prev);
// Clamp in case a pathological amount of duplicate experts in one token exceeded store_capacity above;
// this keeps the read loop below memory-safe (rather than reading uninitialized shared memory).
it_compact = min(it_compact, store_capacity);
for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
const mm_ids_helper_store store_it = store[itc];
const int it = store_it.it();
@@ -134,12 +166,14 @@ static void launch_mm_ids_helper(
const dim3 num_blocks(n_experts, 1, 1);
const dim3 block_size(warp_size, 1, 1);
const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
// 2x slack (see mm_ids_helper) to tolerate a token selecting the same expert more than once.
const size_t nbytes_shared = 2*(size_t) n_tokens*sizeof(mm_ids_helper_store);
GGML_ASSERT(nbytes_shared <= smpbo);
mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
(ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
}
void ggml_cuda_launch_mm_ids_helper(
const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {