fix

Author: shihaobai

lightllm/common/fused_moe/deepep_scatter_gather.py

@@ -217,10 +217,11 @@ def ep_gather(
     input_index: torch.Tensor,
     output_tensor: torch.Tensor,
 ):
-    BLOCK_D = 128  # block size of quantization
-    num_warps = 4
+    BLOCK_D = 1024  # block size of quantization
+    num_warps = 2
     num_tokens = output_tensor.shape[0]
     hidden_size = input_tensor.shape[1]
+    assert hidden_size % BLOCK_D == 0
     grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
     _fwd_kernel_ep_gather[grid](
         num_tokens,