[MOE] try to optimize cu kernel single block execution - distribute cumsum workload from thread 0 to other threads

[yiakwy-xpu-ml-framework-team]

Motivation

Current cuda kernel runs best when work load fitted into a single block execution (when workload large enough triton has better performance since data is more evenly distributed to many blocks, and less possibility for two threads in warp accessing same bank by expert_id )

Before adapting the kernel to the multi-block exeuction. Existing kernel computes expert cumsum in thread 0. Cumsum is a actually very simliar to precan:

For expert number (256 ) In previous single block execution, thread 0 (warp 0) process 256 elements; while in the load balance version, thread 15 has 16 + 16 = 32 elements to process. Hence wavefront is faster to reach.

Modifications

• reduce cusum direct access , not too much positive effects, perhaps expert_num is too small to have influence [x]
• load balance cumsum computation , ~25% boost for small batches; **no effect **for large batches (where existing triton implementation outperforms cuda)

old moe align kernel

new moe align kernel

Test Env

• A100 GPU
• Latest SGLang CUDA Env

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.

Next Step

• (algorithm) Enable efficient multiple blocks execution and adjust threads blocks to large work loads
• enable cuRadix with cub (suported by @BBuf )

[BBuf]

Why is there such a large difference in the Triton benchmark results in the two screenshots above? On the same GPU, Triton's kernel latency should be relatively consistent. @yiakwy-xpu-ml-framework-team

[BBuf]

Based on the proportions in the screenshots above, pr's optimization seems to be slower than main branch ? For example, at bs=1 and seq=64, the CUDA implementation on the main branch takes approximately 26/203 = 0.128 of the time compared to Triton, but in your PR's test results, it shows 20.1/129.79 = 0.154. I'm not sure if you're running the same script on a same GPU. Additionally, the newly added code significantly reduces the readability of the original code. I believe that for small batch_sizes or small seq_length, we don't need to do any optimization. Instead, we should focus on increasing parallelism by using multiple blocks when dealing with large batch sizes or large seq_length now.

[BBuf]

@yiakwy-xpu-ml-framework-team However, your fix for the bug in the benchmark script is correct, and you can create a new PR for the script changes, I can merge it.

[yiakwy-xpu-ml-framework-team]

@BBuf yes triton differences caused by triton 3.2.0 (built from hand for debug); I have udpated result

old

new

and new version is indeed faster 25% at small bathes workloads. And I agree with you enabling multi-blocks efficient execution is more important. It is workign in progress.