improve ep_gather

Author: wanghao7

lightllm/common/fused_moe/deepep_scatter_gather.py

@@ -169,78 +169,84 @@ def ep_scatter(
 
 @triton.jit
 def _fwd_kernel_ep_gather(
-    total_token_num,
     input_tensor,
     input_tensor_stride0,
     input_tensor_stride1,
-    recv_topk_ids,
-    recv_topk_ids_stride0,
-    recv_topk_ids_stride1,
+    recv_topk,
+    recv_topk_stride0,
+    recv_topk_stride1,
     recv_topk_weight,
     recv_topk_weight_stride0,
     recv_topk_weight_stride1,
     input_index,
     input_index_stride0,
     input_index_stride1,
+    expert_start_loc,
     output_tensor,
     output_tensor_stride0,
     output_tensor_stride1,
-    topk_num: tl.constexpr,
+    topk_col: tl.constexpr,
     BLOCK_D: tl.constexpr,
+    HIDDEN_SIZE: tl.constexpr,
+    HIDDEN_SIZE_PAD: tl.constexpr,
 ):
-    cur_block = tl.program_id(0)
-    start_cur_token = tl.program_id(1)
-    grid_num = tl.num_programs(1)
-
-    for cur_token in range(start_cur_token, total_token_num, grid_num):
-        off_d = tl.arange(0, BLOCK_D)
-        accumulator = tl.zeros([BLOCK_D], dtype=tl.float32)
-        for topk_index in range(0, topk_num):
-            expert_id = tl.load(recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index)
-            if expert_id >= 0:
-                source_token_index = tl.load(input_index + cur_token * input_index_stride0 + topk_index)
-                acc_weight = tl.load(recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index)
-                tmp = tl.load(input_tensor + source_token_index * input_tensor_stride0 + cur_block * BLOCK_D + off_d)
-                accumulator += tmp.to(tl.float32) * acc_weight
-
-        tl.store(
-            output_tensor + cur_token * output_tensor_stride0 + cur_block * BLOCK_D + off_d,
-            accumulator.to(output_tensor.dtype.element_ty),
-        )
+    token_id = tl.program_id(0)
+    offset = tl.arange(0, HIDDEN_SIZE_PAD)
+    mask = offset < HIDDEN_SIZE
+    accumulator = tl.zeros([HIDDEN_SIZE_PAD], dtype=tl.float32)
+
+    for start_topk in range(0, topk_col):
+        cur_expert = tl.load(recv_topk + token_id * recv_topk_stride0 + start_topk)
+        if cur_expert >= 0:
+            start_ = tl.load(expert_start_loc + cur_expert)
+            dst = tl.load(input_index + token_id * input_index_stride0 + start_topk) + start_
+
+            weight = tl.load(recv_topk_weight + token_id * recv_topk_weight_stride0 + start_topk)
+            tmp = tl.load(input_tensor + dst + offset)
+            accumulator += tmp.to(tl.float32) * weight
+            
+    tl.store(
+        output_tensor + token_id * output_tensor_stride0 + offset,
+        accumulator.to(output_tensor.dtype.element_ty),
+    )
 
 
 @torch.no_grad()
 def ep_gather(
     input_tensor: torch.Tensor,
-    recv_topk_ids: torch.Tensor,
+    recv_topk: torch.Tensor,
     recv_topk_weight: torch.Tensor,
     input_index: torch.Tensor,
+    expert_start_loc: torch.Tensor,
     output_tensor: torch.Tensor,
 ):
     BLOCK_D = 128  # block size of quantization
     num_warps = 4
     num_tokens = output_tensor.shape[0]
     hidden_size = input_tensor.shape[1]
-    grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
-    _fwd_kernel_ep_gather[grid](
-        num_tokens,
+    grid = min(num_tokens, 65535)
+
+    _fwd_kernel_ep_gather[(grid,)](
         input_tensor,
         input_tensor.stride(0),
         input_tensor.stride(1),
-        recv_topk_ids,
-        recv_topk_ids.stride(0),
-        recv_topk_ids.stride(1),
+        recv_topk,
+        recv_topk.stride(0),
+        recv_topk.stride(1),
         recv_topk_weight,
         recv_topk_weight.stride(0),
         recv_topk_weight.stride(1),
         input_index,
         input_index.stride(0),
         input_index.stride(1),
+        expert_start_loc,
         output_tensor,
         output_tensor.stride(0),
         output_tensor.stride(1),
-        topk_num=recv_topk_ids.shape[1],
+        topk_col=recv_topk.shape[1],
         num_warps=num_warps,
         BLOCK_D=BLOCK_D,
+        HIDDEN_SIZE=hidden_size,
+        HIDDEN_SIZE_PAD = triton.next_power_of_2(hidden_size),
     )
     return