From c5dfaee98aea0d2ac8941174155594bc5a1d31e3 Mon Sep 17 00:00:00 2001 From: project6 Date: Mon, 10 Aug 2026 08:03:08 +0000 Subject: [PATCH] =?UTF-8?q?fix(MoE):=20rewrite=20topk=20kernel=20=E2=80=94?= =?UTF-8?q?=201=20block/row,=20shared=20mem,=20warp-agnostic?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: BI-V100 warp size may be 64 (not 32). Old kernel used dim3(32,4) assuming 4 independent warps per block, but with warpSize=64 two rows shared the same warp → __shfl_sync mixed their data. Debug proof: Row 0 == Row 1, Row 2 == Row 3 (identical outputs). Even rows correct, odd rows duplicated. Fix: 1 block = 1 row = 64 threads (1 per expert). All reductions use shared memory (block_reduce_max/sum/argmax) instead of warp shuffle. Zero warp-size dependency. --- ex_engine/csrc/moe_topk_softmax_v3.cu | 155 +++++++++++++------------- 1 file changed, 75 insertions(+), 80 deletions(-) diff --git a/ex_engine/csrc/moe_topk_softmax_v3.cu b/ex_engine/csrc/moe_topk_softmax_v3.cu index baf9f6e5..99dfe2b3 100644 --- a/ex_engine/csrc/moe_topk_softmax_v3.cu +++ b/ex_engine/csrc/moe_topk_softmax_v3.cu @@ -1,34 +1,56 @@ // moe_topk_softmax_v3.cu — Fused softmax+topk for Qwen3.5 MoE routing // -// VERIFIED on BI-V100 (ivcore10) 2026-08-10: -// weights sum=1.0, no NaN, no duplicate ids, 881 tokens batch OK -// Compiler: corex clang/16, --cuda-gpu-arch=ivcore10 +// 64 experts, topk=8, one block per row, warp shuffle reduction. +// BI-V100 safe: no warp-size assumption (works with warpSize=32 or 64). // -// 64 experts, topk=8, warp shuffle only, zero shared memory -// Each warp handles one token row: 32 threads × 2 values = 64 experts -// -// Based on: TRT-LLM/vllm topk_softmax_kernels + xllm moe_topk_softmax_kernels.cuh +// Each block = 64 threads, each thread owns 1 expert value. +// Softmax: parallel exp + warp reduce. TopK: iterative argmax + mask. #include #include #include -#include - -__device__ __forceinline__ float warp_reduce_max(float val) { - for (int offset = 16; offset > 0; offset >>= 1) - val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); - return val; -} -__device__ __forceinline__ float warp_reduce_sum(float val) { - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} static constexpr int NUM_EXPERTS = 64; -static constexpr int VPT = 2; -static constexpr int THREADS_PER_ROW = NUM_EXPERTS / VPT; -static constexpr int WARPS_PER_CTA = 4; -static constexpr int ROWS_PER_CTA = WARPS_PER_CTA; +static constexpr int BLOCK_SIZE = 64; // 1 thread per expert, 1 block per row + +// Reduce over all 64 threads using shared memory (warp-size agnostic) +__device__ float block_reduce_max(float val, float* smem) { + int tid = threadIdx.x; + smem[tid] = val; + __syncthreads(); + for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]); + __syncthreads(); + } + return smem[0]; +} + +__device__ float block_reduce_sum(float val, float* smem) { + int tid = threadIdx.x; + smem[tid] = val; + __syncthreads(); + for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) smem[tid] += smem[tid + s]; + __syncthreads(); + } + return smem[0]; +} + +// Find global argmax: returns (max_val, max_idx) via shared memory +__device__ void block_argmax(float val, int idx, float* s_val, int* s_idx) { + int tid = threadIdx.x; + s_val[tid] = val; + s_idx[tid] = idx; + __syncthreads(); + for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + if (s_val[tid + s] > s_val[tid]) { + s_val[tid] = s_val[tid + s]; + s_idx[tid] = s_idx[tid + s]; + } + } + __syncthreads(); + } +} __global__ void topk_gating_softmax_kernel( const float* __restrict__ input, @@ -37,61 +59,38 @@ __global__ void topk_gating_softmax_kernel( int32_t* __restrict__ output_source_rows, int num_tokens, int k, bool renormalize ) { - const int row = blockIdx.x * ROWS_PER_CTA + threadIdx.y; + int row = blockIdx.x; if (row >= num_tokens) return; + int tid = threadIdx.x; // 0..63, one per expert - const int tid = threadIdx.x; - const float* row_input = input + row * NUM_EXPERTS; + __shared__ float smem[BLOCK_SIZE]; + __shared__ int smem_idx[BLOCK_SIZE]; - float vals[VPT]; - int my_indices[VPT]; - #pragma unroll - for (int i = 0; i < VPT; i++) { - int col = tid * VPT + i; - vals[i] = row_input[col]; - my_indices[i] = col; - } + // Load gating logit for this expert + float val = input[row * NUM_EXPERTS + tid]; - float tmax = vals[0]; - for (int i = 1; i < VPT; i++) tmax = fmaxf(tmax, vals[i]); - float row_max = warp_reduce_max(tmax); - - float tsum = 0.0f; - #pragma unroll - for (int i = 0; i < VPT; i++) { - vals[i] = expf(vals[i] - row_max); - tsum += vals[i]; - } - float row_sum = warp_reduce_sum(tsum); - float inv_sum = 1.0f / row_sum; - - #pragma unroll - for (int i = 0; i < VPT; i++) vals[i] *= inv_sum; + // Softmax: max-subtract, exp, normalize + float row_max = block_reduce_max(val, smem); + val = expf(val - row_max); + float row_sum = block_reduce_sum(val, smem); + val *= (1.0f / row_sum); + // Output pointers for this row float* out_w = output_weights + row * k; int32_t* out_idx = output_indices + row * k; int32_t* out_src = output_source_rows + row * k; + // Iterative top-k: find max, write, mask, repeat float topk_sum = 0.0f; + float my_val = val; // will be set to -1 when selected for (int ki = 0; ki < k; ki++) { - float local_max = -1.0f; - int local_idx = -1; - #pragma unroll - for (int i = 0; i < VPT; i++) { - if (vals[i] > local_max) { - local_max = vals[i]; - local_idx = my_indices[i]; - } - } - - float global_max = warp_reduce_max(local_max); - bool is_winner = (local_max == global_max && local_max > 0.0f); - unsigned winner_mask = __ballot_sync(0xFFFFFFFF, is_winner); - int first_winner = __ffs(winner_mask) - 1; - - float winner_val = __shfl_sync(0xFFFFFFFF, local_max, first_winner); - int winner_idx = __shfl_sync(0xFFFFFFFF, local_idx, first_winner); + block_argmax(my_val, tid, smem, smem_idx); + // Thread 0 has the winner + float winner_val = smem[0]; + int winner_idx = smem_idx[0]; + // Broadcast via shared memory (already in smem[0]) + __syncthreads(); if (tid == 0) { out_w[ki] = winner_val; @@ -100,17 +99,15 @@ __global__ void topk_gating_softmax_kernel( } topk_sum += winner_val; - #pragma unroll - for (int i = 0; i < VPT; i++) { - if (my_indices[i] == winner_idx) - vals[i] = -1.0f; - } + // Mask out the selected expert + if (tid == winner_idx) my_val = -1.0f; + __syncthreads(); } if (renormalize && tid == 0) { - float inv_topk = 1.0f / (topk_sum + 1e-8f); + float inv = 1.0f / (topk_sum + 1e-8f); for (int ki = 0; ki < k; ki++) - out_w[ki] *= inv_topk; + out_w[ki] *= inv; } } @@ -118,7 +115,8 @@ std::vector moe_topk_softmax( torch::Tensor gating_output, int64_t topk, bool renormalize ) { int num_tokens = gating_output.size(0); - TORCH_CHECK(gating_output.size(1) == 64, "Specialized for 64 experts"); + int num_experts = gating_output.size(1); + TORCH_CHECK(num_experts == 64, "Specialized for 64 experts, got ", num_experts); auto opts_f = torch::dtype(torch::kFloat32).device(gating_output.device()); auto opts_i = torch::dtype(torch::kInt32).device(gating_output.device()); @@ -126,12 +124,9 @@ std::vector moe_topk_softmax( auto topk_ids = torch::empty({num_tokens, topk}, opts_i); auto token_expert_ids = torch::empty({num_tokens, topk}, opts_i); - auto input_f32 = gating_output.to(torch::kFloat32); + auto input_f32 = gating_output.to(torch::kFloat32).contiguous(); - dim3 block(THREADS_PER_ROW, WARPS_PER_CTA); - dim3 grid((num_tokens + ROWS_PER_CTA - 1) / ROWS_PER_CTA); - - topk_gating_softmax_kernel<<>>( input_f32.data_ptr(), topk_weights.data_ptr(), @@ -144,5 +139,5 @@ std::vector moe_topk_softmax( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("moe_topk_softmax", &moe_topk_softmax, - "Fused softmax+topk for MoE routing (64 experts, warp shuffle, zero SMEM)"); + "Fused softmax+topk for MoE routing (64 experts, shared mem, warp-agnostic)"); }