Add deepseek_v3 fused gate

[NovTi]

Add deepseek v3 fused gate module

[BBuf]

Please refactor this code into a standalone function, which can be directly used from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/topk.py#L111-L147.

[NovTi]

Do you mean I separate the reference implementation into a standalone function?

[BBuf]

Yeah.

[NovTi]

Got it, I will do that

[BBuf]

Why not directly compare output and output_ref instead of sorting them?

[NovTi]

This is weird, kernel sometimes will output exact same output but in a different order. I checked the following steps and the output order does not matter so I used this way to do the unit test, is this ok?

[BBuf]

We need to determine at which specific step of the fused kernel this inconsistency in order occurs. Additionally, we need to clarify whether running the PyTorch implementation twice with the same input would result in inconsistent output orders. Finally, if you believe that the current order inconsistency does not affect the fused MoE accuracy, you need to provide an end-to-end result, such as running the GSM8K test with the DeepSeek V3 model.

[NovTi]

I see, I will check the inconsistency inside the kernel. I cannot run e2e test on my server, Yineng will help me do the test

[BBuf]

Can you add a benchmark script? Maybe refer to https://github.com/sgl-project/sglang/tree/main/sgl-kernel/benchmark

[NovTi]

Sure

[BBuf]

It seems more appropriate to name it deepseekv3_fused_gate here, as models from the deepseek series can all go through this gate function.

[NovTi]

This is not a generalized kernel, it only works for deepseek v3 671b model

[BBuf]

I see, thanks.

[zhyncs]

I think it also works for DeepSeek V2 VL

[BBuf]

Synchronization is not allowed in CUDA kernel's host code, as it will cause CUDA graphs to crash. Can you remove it?

[NovTi]

Thanks, I will update these

[BBuf]

Please add Adapted from https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu#L231

[BBuf]

In TensorRT-LLM, the fused MoE module, in addition to the fused_gate here, also includes the trt_moe_expand_and_permute, the CUTLASS grouped GEMM, and the trt_moe_unpermute_and_reduce processes. Compared to the MoE implemented in Triton, the advantage of TensorRT-LLM's approach is that it does not require padding, which saves some computational overhead, and the CUTLASS implementation may have greater performance potential, especially on Hopper architecture. I conducted an experiment at https://github.com/sgl-project/sglang/tree/bbuf_tmp, where I successfully connected trt_moe_expand_and_permute, trt_moe_unpermute_and_reduce, and FlashInfer's grouped GEMM in sgl-kernel to run correctness comparisons with the Triton fused MoE operator in the case of bfloat16 dtype. However, it seems that the current performance is still significantly worse than Triton's. This could be due to performance issues with FlashInfer's grouped GEMM on specific shapes. Additionally, FlashInfer's GEMM does not currently support scaled FP8 or INT8 GEMM. If anyone is interested, we can discuss whether to directly integrate TensorRT's fused MoE as a backend into sglang or to use FlashInfer's approach, which would require a customization of FlashInfer for grouped GEMM. cc @zhyncs

[zhyncs]

directly integrate TensorRT's fused MoE as a backend into sglang

sounds good @BBuf

[BBuf]

directly integrate TensorRT's fused MoE as a backend into sglang

sounds good @BBuf

Yeah, I can have a try.