diff --git a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu index 70738390b..ef7495ccf 100644 --- a/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu +++ b/python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu @@ -1,7 +1,6 @@ #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; @@ -9,17 +8,12 @@ struct mm_ids_helper_store { 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 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,11 +26,11 @@ 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; @@ -44,73 +38,106 @@ 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: 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(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::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(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(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(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(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(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::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(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 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, 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, 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<<>> - (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse); + mm_ids_helper<<>> + (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( diff --git a/python/freetoken/kernel/ggml.py b/python/freetoken/kernel/ggml.py index c9960e594..d18998b81 100644 --- a/python/freetoken/kernel/ggml.py +++ b/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. """ diff --git a/tests/kernels/test_ggml.py b/tests/kernels/test_ggml.py index 9db4aaf6b..375bcfc92 100644 --- a/tests/kernels/test_ggml.py +++ b/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))