fix(moe): preserve duplicate gguf expert routes

python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu

@@ -1,25 +1,19 @@
 #include "common.cuh"
 #include "mmid.cuh"
 
-// To reduce shared memory use, store "it" and "iex_used" with 22/10 bits each.
 struct mm_ids_helper_store {
     uint32_t data;
 
     __device__ mm_ids_helper_store(const uint32_t it, const uint32_t iex_used) {
         data = (it & 0x003FFFFF) | (iex_used << 22);
     }
 
-    __device__ uint32_t it() const {
-        return data & 0x003FFFFF;
-    }
-
-    __device__ uint32_t iex_used() const {
-        return data >> 22;
-    }
+    __device__ uint32_t it() const { return data & 0x003FFFFF; }
+    __device__ uint32_t iex_used() const { return data >> 22; }
 };
 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
+// 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; };
@@ -32,85 +26,118 @@ 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) {
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1,
+        const bool write_inverse, const bool use_shared) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
     const int expert = blockIdx.x;
-
     extern __shared__ char data_mm_ids_helper[];
     mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
 
     int nex_prev   = 0; // Number of columns for experts with a lower index.
     int it_compact = 0; // Running index for the compact slice of this expert.
 
     if constexpr (n_expert_used_template == 0) {
-        // Generic implementation:
         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];
                 nex_prev += (unsigned) expert_used < (unsigned) expert; // FreeToken: unsigned, so an inactive route (negative id) sorts after every expert
-                if (expert_used == expert) {
-                    iex_used = iex;
-                }
-            }
-
-            if (iex_used != -1) {
-                store[it_compact] = mm_ids_helper_store(it, iex_used);
-            }
-
-            if (warp_reduce_any<warp_size>(iex_used != -1)) {
-                it_compact++;
             }
         }
     } else {
-        // Implementation optimized for specific numbers of experts used:
-        // a warp holds a whole number of token slots, so the slot count is padded to a power of 2
         constexpr int neu_padded = mm_ids_pow2<n_expert_used_template>::value;
         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;
-
-            const int iex = threadIdx.x % neu_padded; // The index at which the expert is used, if any.
+            const int iex = threadIdx.x % neu_padded;
             const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
                 ids[it*si1 + iex] : INT_MAX;
-            const int iex_used = expert_used == expert ? iex : -1;
             nex_prev += (unsigned) expert_used < (unsigned) expert; // FreeToken: unsigned, so an inactive route (negative id) sorts after every 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);
-
-            // 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;
-#pragma unroll
-            for (int offset = neu_padded; offset < warp_size; offset += neu_padded) {
-                const int tmp = __shfl_up_sync(0xFFFFFFFF, it_compact_add_self, offset, warp_size);
-                if (threadIdx.x >= static_cast<unsigned int>(offset)) {
-                    it_compact_add_lower += tmp;
+            if (use_shared) {
+                const int match = expert_used == expert;
+                int rank = match;
+                for (int offset = 1; offset < warp_size; offset *= 2) {
+                    const int prior = __shfl_up_sync(0xFFFFFFFF, rank, offset, warp_size);
+                    if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                        rank += prior;
+                    }
                 }
+                const int count = __shfl_sync(0xFFFFFFFF, rank, warp_size - 1, warp_size);
+                if (match) {
+                    store[it_compact + rank - 1] = mm_ids_helper_store(it, iex);
+                }
+                it_compact += count;
             }
-
-            if (iex_used != -1) {
-                store[it_compact + it_compact_add_lower] = 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:
-            it_compact += __shfl_sync(0xFFFFFFFF, it_compact_add_lower + it_compact_add_self, warp_size - 1, warp_size);
         }
     }
     nex_prev = warp_reduce_sum<warp_size>(nex_prev);
-    ggml_cuda_syncwarp();
-
-    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();
-        const int iex_used = store_it.iex_used();
-        ids_dst[nex_prev + itc] = it*n_expert_used + iex_used;
-        // ids_src1 holds the forward map, or the inverse map (token slot -> compact row) for quant dedup
-        if (write_inverse) {
-            ids_src1[it*n_expert_used + iex_used] = nex_prev + itc;
+    if (use_shared) {
+        ggml_cuda_syncwarp();
+        for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
+            const mm_ids_helper_store route = store[itc];
+            const int it = route.it();
+            const int iex = route.iex_used();
+            const int dst = nex_prev + itc;
+            ids_dst[dst] = it*n_expert_used + iex;
+            if (write_inverse) {
+                ids_src1[it*n_expert_used + iex] = dst;
+            } else {
+                ids_src1[dst] = it*sis1 + iex % nchannels_y;
+            }
+        }
+    } else {
+        // Count first so every route can write directly to its final expert slice.
+        // A token may use the same expert more than once; a per-token compact row loses routes.
+        if constexpr (n_expert_used_template == 0) {
+            for (int it = 0; it < n_tokens; ++it) {
+                for (int iex0 = 0; iex0 < n_expert_used; iex0 += warp_size) {
+                    const int iex = iex0 + threadIdx.x;
+                    const int match = iex < n_expert_used && ids[it*si1 + iex] == expert;
+                    int rank = match;
+                    for (int offset = 1; offset < warp_size; offset *= 2) {
+                        const int prior = __shfl_up_sync(0xFFFFFFFF, rank, offset, warp_size);
+                        if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                            rank += prior;
+                        }
+                    }
+                    const int count = __shfl_sync(0xFFFFFFFF, rank, warp_size - 1, warp_size);
+                    if (match) {
+                        const int dst = nex_prev + it_compact + rank - 1;
+                        ids_dst[dst] = it*n_expert_used + iex;
+                        if (write_inverse) {
+                            ids_src1[it*n_expert_used + iex] = dst;
+                        } else {
+                            ids_src1[dst] = it*sis1 + iex % nchannels_y;
+                        }
+                    }
+                    it_compact += count;
+                }
+            }
         } else {
-            ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
+            constexpr int neu_padded = mm_ids_pow2<n_expert_used_template>::value;
+            for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
+                const int it = it0 + threadIdx.x / neu_padded;
+                const int iex = threadIdx.x % neu_padded;
+                const int match = iex < n_expert_used && it < n_tokens && ids[it*si1 + iex] == expert;
+                int rank = match;
+                for (int offset = 1; offset < warp_size; offset *= 2) {
+                    const int prior = __shfl_up_sync(0xFFFFFFFF, rank, offset, warp_size);
+                    if (threadIdx.x >= static_cast<unsigned int>(offset)) {
+                        rank += prior;
+                    }
+                }
+                const int count = __shfl_sync(0xFFFFFFFF, rank, warp_size - 1, warp_size);
+                if (match) {
+                    const int dst = nex_prev + it_compact + rank - 1;
+                    ids_dst[dst] = it*n_expert_used + iex;
+                    if (write_inverse) {
+                        ids_src1[it*n_expert_used + iex] = dst;
+                    } else {
+                        ids_src1[dst] = it*sis1 + iex % nchannels_y;
+                    }
+                }
+                it_compact += count;
+            }
         }
     }
 
@@ -131,20 +158,19 @@ template <int n_expert_used_template>
 static void 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_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
-    GGML_ASSERT(n_tokens          < (1 << 22) && "too few bits in mm_ids_helper_store");
-    GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
-
     const int id = ggml_cuda_get_device();
     const int warp_size = ggml_cuda_info().devices[id].warp_size;
     const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
-    CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);
-
+    const size_t needed_shared = (size_t) n_tokens*n_expert_used_var*sizeof(mm_ids_helper_store);
+    const bool use_shared = n_expert_used_template != 0 && n_tokens < (1 << 22) &&
+                            n_expert_used_var < (1 << 10) && needed_shared <= smpbo;
+    if (use_shared) {
+        CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);
+    }
     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);
-    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);
+    mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, use_shared ? needed_shared : 0, stream>>>
+        (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse, use_shared);
 }
 
 void ggml_cuda_launch_mm_ids_helper(

python/freetoken/kernel/ggml.py

@@ -100,6 +100,7 @@ def mul_mat_id(
 ) -> torch.Tensor:
     """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``.
 
+    Repeated nonnegative expert IDs within a token are allowed; each route produces its own output.
     ``scale`` is an fp32 ``[E]`` view (strided is fine) that multiplies each NVFP4 expert's rows; it needs ``T <= mmvq_max_tokens``.
     ``inactive`` says ``ids`` may hold -1 (hybrid decode's CPU routes): those rows come back zero. Above ``mmvq_max_tokens`` it costs a memset and a zeroing pass, so pass it only then.
     """

tests/kernels/test_ggml.py

@@ -123,6 +123,18 @@ def test_iq1_m_matrix_and_moe_routes_match_cpu(dtype, batch):
     torch.testing.assert_close(actual.cpu(), expected, atol=2e-4, rtol=1e-3)
 
 
+@pytest.mark.parametrize("qtype", [QType.Q4_0, QType.IQ1_M])
+@pytest.mark.parametrize("top_k", [3, 8])
+def test_ggml_mmq_duplicate_routes_match_cpu(qtype, top_k):
+    k = 512 if qtype == QType.IQ1_M else 256
+    w, ref = _iq1_m_weights((4, 64, k), True) if qtype == QType.IQ1_M else _q4_weights((4, 64, k))
+    x = _activation((17, top_k, k), torch.bfloat16)
+    ids = torch.tensor([[t % 4] * top_k for t in range(17)], dtype=torch.int32, device="cuda")
+    actual = ggml.mul_mat_id(w, int(qtype), x, ids)
+    expected = torch.einsum("tjk,tjok->tjo", x.cpu().float(), ref[ids.cpu().long()]).to(x.dtype)
+    torch.testing.assert_close(actual.cpu(), expected, atol=2e-4, rtol=1e-3)
+
+
 @pytest.mark.parametrize("dtype", DTYPES)
 def test_q4_shared_q8_matches_internal_quantization(dtype):
     w, _ = _q4_weights((64, 256))