fix it

ggml/src/ggml-cuda/mmid.cu

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

tests/test-backend-ops.cpp

@@ -4404,7 +4404,7 @@ struct test_mul_mat_hadamard : public test_mul_mat {
     }
 };
 
-static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
+static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, bool dup_ids = false) {
     std::random_device rd;
     std::default_random_engine rng(rd());
     for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
@@ -4417,6 +4417,12 @@ static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
                     data[i] = i % n_mats;
                 }
                 std::shuffle(data.begin(), data.end(), rng);
+                if (dup_ids && t->ne[0] >= 2) {
+                    // some models (e.g. expert-pruned MoE checkpoints) can end up with a
+                    // routing table that selects the same expert twice for one token,
+                    // make sure this is handled correctly (see issue #26588)
+                    data[1] = data[0];
+                }
                 ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
             }
         } else {
@@ -4435,9 +4441,10 @@ struct test_mul_mat_id : public test_case {
     const int64_t m;
     const int64_t n;
     const int64_t k;
+    const bool dup_ids; // force a duplicated expert id within a token's selected experts
 
     std::string vars() override {
-        return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k);
+        return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, dup_ids);
     }
 
     double max_nmse_err() override {
@@ -4459,10 +4466,11 @@ struct test_mul_mat_id : public test_case {
 
     test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
             int n_mats = 8, int n_used = 2, bool b = false,
-            int64_t m = 32, int64_t n = 32, int64_t k = 32)
+            int64_t m = 32, int64_t n = 32, int64_t k = 32, bool dup_ids = false)
         : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
-            m(m), n(n), k(k) {
+            m(m), n(n), k(k), dup_ids(dup_ids) {
             GGML_ASSERT(n_used <= n_mats);
+            GGML_ASSERT(!dup_ids || n_used >= 2);
         }
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
@@ -4487,7 +4495,7 @@ struct test_mul_mat_id : public test_case {
     }
 
     void initialize_tensors(ggml_context * ctx) override {
-        init_mul_mat_id_tensors(ctx, n_mats);
+        init_mul_mat_id_tensors(ctx, n_mats, dup_ids);
     }
 };
 
@@ -9029,7 +9037,16 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880));
     test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880));
 
+    // a malformed routing table can select the same expert twice for one token (see issue #26588),
+    // make sure quantized MMQ-style paths do not assume unique expert ids per token
+    for (ggml_type type_a : {GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_MXFP4}) {
+        test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, 6, false, 128, 1, 256, true));
+        test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, 6, false, 128, 32, 256, true));
+        test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880, true));
+    }
+
     for (ggml_type type_a : all_types) {
+
         test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 2, false, 64, 16, 3*ggml_blck_size(type_a)));
     }