Add deepseek style fused moe group gate selection kernel

[qingquansong]

Motivation

PR adapted and improved from #3191
Rewrite Macro. Extended to support all power of 2 # expert & # expert group, also all # topk_group & # topk use cases + dtype support fp16/bf16/fp32.

TODO:

• Add trt llm reference? (currently cannot find the original reference)
• e2e perf run

NOTE:

• Current CUDA kernel matches the eager mode results for supported cases, but torch compile with static or dynamic will cause original torch implementation results mismatch. (not a problem for this pr but worth investigating in the future)
• Currently only support # experts is power of 2
• Currently force # experts / # expert group <= 32 as we fix size of AlignedArray in expression (MAX_VPT=32 later we can make this dynamic equal to the params.VPT size can improve the speed for smaller cases)

Test

Unit tests (in pr)
Speed test for deepseek v3 case: 256 expert + 8 expert group + first select top4 expert group + then top8 final selected expert

seq_length	Original (eager)	Original (compile static)	Original (compile dynamic)	SGL Kernel (cutlass dtype/array)	SGL Kernel (cutlass fp16)	SGL Kernel (cutlass fp32)	SGL Kernel (bfloat16 native dtype/array)
5k	251.12	199.94	213.01	26.40	26.56	36.06	31.26
10k	445.47	381.00	399.78	37.82	37.22	52.93	44.67
15k	589.89	510.21	534.14	43.39	44.74	74.46	52.54
20k	772.83	678.59	707.28	54.11	54.50	87.46	64.77
25k	957.54	848.14	879.70	59.84	59.30	100.05	70.98
30k	1100.74	979.17	1015.07	68.10	68.70	118.94	81.60
35k	1284.54	1152.86	1190.45	78.50	78.94	137.25	93.86
40k	1432.72	1283.01	1325.00	84.38	86.05	152.22	100.64

Modifications

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]

@qingquansong Great job! Would you like to consider of usage AlignedArray, and native datatype? Of course I can create amd/ck/fused_moe_gate.cu counter part, but it will be great if I can reuse the codes.

I think the fusion algorithm is great and we can do some deeper engineering work later.

[yiakwy-xpu-ml-framework-team]

Hi @qingquansong TRT_LLM has removed "tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h" from v0.16.0 to
"tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h" in v0.17.0 and hide the implementation a static library.

Could you tell me moe_fused_gate_impl done by our team without referencing any implementation before (#3191) ? (copyright issue)

If that is not case, that is great! Do we have ncu profiling for sharing ?

[qingquansong]

Hey @yiakwy-xpu-ml-framework-team Thanks! I seem to not able to find it from the file https://github.com/NVIDIA/TensorRT-LLM/blob/v0.16.0/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu but only our previous moeSoftmax kernel is adapted from there, maybe this is a new implementation or just referred to some similar code to implement? @BBuf do you happen to find it somewhere in the old trt version so we can put as reference? Thanks both!

[qingquansong]

@qingquansong Great job! Would you like to consider of usage Aligned job, and native datatype? Of course I can create amd/ck/fused_moe_gate.cu counter part, but it will be great if I can reuse the codes.

I think the fusion algorithm is great and we can do some deeper engineering work later.

Definitely, pushed one version with a USE_ROCM switch of the native type and the cutlass one and let me know if my understanding is correct. The speed results are put in the description for reference. Thank you!

[qingquansong]

The PR is picked up at #4530 cc @zhyncs @yiakwy-xpu-ml-framework-team @BBuf @hebiao064 @zcnrex @HandH1998