Add moe topk softmax templated from vllm

[qingquansong]

Motivation

#2965

Modifications

1. Cherry picked current vllm MoE topk softmax kernel template (with a fix on naming typo for token_expert_indices)
2. Polish util func warpReduceMax / blockReduceMax for handle AMD use case as well.

Tests

Unit tests + benchmarking aligned with vllm counterpart

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[yiakwy-xpu-ml-framework-team]

Hi @qingquansong

once this PR #4432 is merged

#ifndef USE_ROCM
  #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \
    __shfl_xor_sync(uint32_t(-1), var, lane_mask)
  #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
    __shfl_xor_sync(uint32_t(-1), var, lane_mask, width)
#else
  #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask)
  #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
    __shfl_xor(var, lane_mask, width)
#endif

you can use __shfl_xor_sync directly. hence no need to have these lines.

Does it sound good to you ?

[yiakwy-xpu-ml-framework-team]

The modification to these function will be revert in #4432

cc @zhyncs

[yiakwy-xpu-ml-framework-team]

I will keep these lines since you have the ROCM specific macro (since CUDA operation is no longer safe if we employ this approach) in many places. But #4432 is merged. The macro is no longer needed.

[qingquansong]

Thank you! I'll change those and remove the definition.

[yiakwy-xpu-ml-framework-team]

hipCUB is experimental one, we can try it. But it introduces new dependencies.

Could just use some simple reduction kernel ?

[qingquansong]

Definitely. We can change to customized reductions for both max and sum. I'll do it together with the macro change in a follow-up pr. How about the following one? I can test the correctness on CUDA and may need your help for AMD machine testing.

__device__ __forceinline__ float warpReduceSum(float sum_value) {
  sum_value += __shfl_xor_sync(0xffffffff, sum_value, 16);
  sum_value += __shfl_xor_sync(0xffffffff, sum_value, 8);
  sum_value += __shfl_xor_sync(0xffffffff, sum_value, 4);
  sum_value += __shfl_xor_sync(0xffffffff, sum_value, 2);
  sum_value += __shfl_xor_sync(0xffffffff, sum_value, 1);
  return sum_value;
}

__device__ __forceinline__ float blockReduceSum(float sum_value) {
  static __shared__ float warpLevelSums[WARP_SIZE];
  const int laneId = threadIdx.x % WARP_SIZE;
  const int warpId = threadIdx.x / WARP_SIZE;

  sum_value = warpReduceSum(sum_value);

  if (laneId == 0) warpLevelSums[warpId] = sum_value;
  __syncthreads();

  sum_value = (threadIdx.x < blockDim.x / WARP_SIZE) ? warpLevelSums[laneId] : 0;
  if (warpId == 0) sum_value = warpReduceSum(sum_value);

  return sum_value;
}

[yiakwy-xpu-ml-framework-team]

resovled in #4448

[yiakwy-xpu-ml-framework-team]

But also I recommend to use shlf_xor based implementation. The old solution from fasterTransformer(later incorporated into TRT-LLM) uses heavily shared memory for reduction:

const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());

WIth shlf_xor based implementation, then you can get better result.

@qingquansong

[qingquansong]

Hi @qingquansong

once this PR #4432 is merged

#ifndef USE_ROCM
  #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \
    __shfl_xor_sync(uint32_t(-1), var, lane_mask)
  #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
    __shfl_xor_sync(uint32_t(-1), var, lane_mask, width)
#else
  #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask)
  #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
    __shfl_xor(var, lane_mask, width)
#endif

you can use __shfl_xor_sync directly. hence no need to have these lines.

Does it sound good to you ?

Sounds great! I'll remove the marco definition and change back to use __shfl_xor_sync directly