update gather

Author: shihaobai

lightllm/common/fused_moe/deepep_scatter_gather.py

@@ -185,27 +185,31 @@ def _fwd_kernel_ep_gather(
     output_tensor,
     output_tensor_stride0,
     output_tensor_stride1,
+    HIDDEN_SIZE: tl.constexpr,
+    HIDDEN_SIZE_PAD: tl.constexpr,
     topk_num: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
-    cur_block = tl.program_id(0)
-    start_cur_token = tl.program_id(1)
-    grid_num = tl.num_programs(1)
+    start_cur_token = tl.program_id(0)
+    grid_num = tl.num_programs(0)
+
+    offset_d = tl.arange(0, HIDDEN_SIZE_PAD)
+    mask = offset_d < HIDDEN_SIZE
 
     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)
+        accumulator = tl.zeros([HIDDEN_SIZE_PAD], 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)
+                tmp = tl.load(input_tensor + source_token_index * input_tensor_stride0 + offset_d, mask=mask)
                 accumulator += tmp.to(tl.float32) * acc_weight
 
         tl.store(
-            output_tensor + cur_token * output_tensor_stride0 + cur_block * BLOCK_D + off_d,
+            output_tensor + cur_token * output_tensor_stride0 + offset_d,
             accumulator.to(output_tensor.dtype.element_ty),
+            mask=mask,
         )
 
 
@@ -221,7 +225,7 @@ def ep_gather(
     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))
+    grid = (min(num_tokens, 32768),)
     _fwd_kernel_ep_gather[grid](
         num_tokens,
         input_tensor,
@@ -240,6 +244,8 @@ def ep_gather(
         output_tensor.stride(0),
         output_tensor.stride(1),
         topk_num=recv_topk_ids.shape[1],
+        HIDDEN_SIZE=hidden_size,
+        HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
         num_warps=num_warps,
         BLOCK_D=BLOCK_D,
     )