From 12d6818380487171faefd66479f086219daf7492 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 2 Oct 2025 21:55:53 +0800 Subject: [PATCH] Tiny fix ep_gather behavior different in CI (#11130) --- python/sglang/srt/layers/moe/ep_moe/kernels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 08660812d..d8e221d5c 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -1104,10 +1104,10 @@ def ep_gather( input_index: torch.Tensor, output_tensor: torch.Tensor, ): - BLOCK_D = 1024 if not is_in_ci() else 128 # block size of quantization num_warps = 2 num_tokens = output_tensor.shape[0] hidden_size = input_tensor.shape[1] + BLOCK_D = 128 if hidden_size % 1024 != 0 else 1024 # block size of quantization assert hidden_size % BLOCK_D == 0 grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024)) _fwd_kernel_ep_gather[grid](