EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Tiny fix ep_gather behavior different in CI (#11130)

Browse Source
This commit is contained in:
fzyzcjy
2025-10-02 21:55:53 +08:00
committed by GitHub
parent b65db0287b
commit 12d6818380
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/moe/ep_moe/kernels.py
View File

@@ -1104,10 +1104,10 @@ def ep_gather(
input_index: torch.Tensor, input_index: torch.Tensor,
output_tensor: torch.Tensor, output_tensor: torch.Tensor,
): ):
BLOCK_D = 1024 if not is_in_ci() else 128 # block size of quantization
num_warps = 2 num_warps = 2
num_tokens = output_tensor.shape[0] num_tokens = output_tensor.shape[0]
hidden_size = input_tensor.shape[1] 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 assert hidden_size % BLOCK_D == 0
grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024)) grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
_fwd_kernel_ep_gather[grid]( _fwd_kernel_ep_gather[grid](