[EP] Add cuda kernel for moe_ep_pre_reorder (#6699)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

sgl-kernel/benchmark/bench_moe_ep_pre_reorder.py

@@ -0,0 +1,100 @@
+import itertools
+
+import torch
+import triton
+import triton.language as tl
+from sgl_kernel import ep_moe_pre_reorder
+
+from sglang.srt.layers.moe.ep_moe.kernels import pre_reorder_triton_kernel
+
+batch_sizes = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]
+configs = [(bs,) for bs in batch_sizes]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size"],
+        x_vals=[list(_) for _ in configs],
+        line_arg="provider",
+        line_vals=["cuda", "triton"],
+        line_names=["CUDA Kernel", "Triton Kernel"],
+        styles=[("green", "-"), ("orange", "-")],
+        ylabel="us",
+        plot_name="ep-moe-pre-reorder-performance",
+        args={},
+    )
+)
+def benchmark(batch_size, provider):
+    dtype = torch.float32
+    device = torch.device("cuda")
+    hidden_size, topk, start_expert_id, end_expert_id, block_size = 4096, 8, 0, 255, 512
+
+    # Allocate fresh tensors for every run to match bench_moe_fused_gate style
+    def alloc_tensors():
+        input_ = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
+        gateup_input = torch.zeros(
+            batch_size * topk, hidden_size, dtype=dtype, device=device
+        )
+        src2dst = torch.randint(
+            0, batch_size * topk, (batch_size, topk), dtype=torch.int32, device=device
+        )
+        topk_ids = torch.randint(
+            start_expert_id,
+            end_expert_id + 1,
+            (batch_size, topk),
+            dtype=torch.int32,
+            device=device,
+        )
+        a1_scales = torch.rand(
+            end_expert_id - start_expert_id + 1, dtype=torch.float32, device=device
+        )
+        return input_, gateup_input, src2dst, topk_ids, a1_scales
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "cuda":
+
+        def run_cuda():
+            inp, gout, s2d, tk_ids, scales = alloc_tensors()
+            ep_moe_pre_reorder(
+                inp,
+                gout,
+                s2d,
+                tk_ids,
+                scales,
+                start_expert_id,
+                end_expert_id,
+                topk,
+                True,
+            )
+
+        ms, min_ms, max_ms = triton.testing.do_bench(run_cuda, quantiles=quantiles)
+
+    elif provider == "triton":
+
+        def run_triton():
+            inp, gout, s2d, tk_ids, scales = alloc_tensors()
+            pre_reorder_triton_kernel[(batch_size,)](
+                inp.view(-1),
+                gout.view(-1),
+                s2d.view(-1),
+                tk_ids.view(-1),
+                scales,
+                start_expert_id,
+                end_expert_id,
+                topk,
+                hidden_size,
+                block_size,
+                True,
+            )
+
+        ms, min_ms, max_ms = triton.testing.do_bench(run_triton, quantiles=quantiles)
+
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)