Fix benchmark_sglang_fused_moe_triton.py (#18940)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py

@@ -18,7 +18,13 @@ from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
 from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
-from sglang.srt.layers.moe.topk import TopK, TopKConfig, select_experts
+from sglang.srt.layers.moe.topk import (
+    TopK,
+    TopKConfig,
+    TopKOutputFormat,
+    select_experts,
+)
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 
 
 def fused_moe_triton_api(
@@ -32,8 +38,8 @@ def fused_moe_triton_api(
         top_k=topk,
         renormalize=False,
         use_grouped_topk=False,
+        output_format=TopKOutputFormat.TRITON_KERNEL,
     )
-    topk_op.use_triton_kernels = True
     triton_topk_output = topk_op.forward_cuda(
         hidden_states=x,
         router_logits=input_gating,
@@ -199,6 +205,10 @@ def main():
     parser.add_argument("--trust-remote-code", action="store_true")
     args = parser.parse_args()
 
+    # Initialize global server args (required by SGLang MoE kernels)
+    server_args = ServerArgs(model_path=args.model)
+    set_global_server_args_for_scheduler(server_args)
+
     try:
         if not torch.distributed.is_initialized():
             torch.distributed.init_process_group(
@@ -217,8 +227,8 @@ def main():
         )
 
         initialize_model_parallel(
-            tensor_model_parallel_size=args.ep_size,
+            tensor_model_parallel_size=1,
-            pipeline_model_parallel_size=args.tp_size,
+            expert_model_parallel_size=1,
         )
 
         model_config = get_model_config(args.model, args.tp_size, args.ep_size)