dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix(MoE): rewrite topk kernel — 1 block/row, shared mem, warp-agnostic

Browse Source 
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.
This commit is contained in:
project6
2026-08-10 08:03:08 +00:00
parent a1fc56d5b0
commit c5dfaee98a
1 changed files with 75 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

155
ex_engine/csrc/moe_topk_softmax_v3.cu
View File

@@ -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 <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
__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<torch::Tensor> 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<torch::Tensor> 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<<<grid, block, 0,
topk_gating_softmax_kernel<<<num_tokens, BLOCK_SIZE, 0,
c10::cuda::getCurrentCUDAStream()>>>(
input_f32.data_ptr<float>(),
topk_weights.data_ptr<float>(),
@@ -144,5 +139,5 @@ std::vector<torch::Tensor> 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)");
}