fix moe_align1 kernel performance issue in prefill stage. (#718)

old version:
test groped fused moe speed.py 200
token num: 200 cost time:0.0011632442474365234 s
.256
token num: 256 cost time: 0.0011243820198429688 S
.8192
token num:8192 cost time: 0.05202174186706543 s
new version:
test groped fused moe speed.py 200
token num: 200 cost time:0.0011744499206542969 5
256
token num:256 cost time:0.0010919570922851562 s
.8192
token num: 8192 cost time: 0.003216266632080078 s

8192 token 10x faster.

Author: hiworldwzj

lightllm/common/fused_moe/grouped_fused_moe.py

@@ -114,29 +114,37 @@ def moe_align1_kernel(
     experts_topk_weight,  # [expert_num, token_num * topk_num]
     experts_topk_weight_stride0,
     experts_topk_weight_stride1,
-    TOKEN_BLOCK_N: tl.constexpr,
+    TOKEN_BLOCK_SIZE: tl.constexpr,
+    NUM_STAGE: tl.constexpr,
 ):
 
     expert_id = tl.program_id(axis=0)
-    n_range = tl.arange(0, TOKEN_BLOCK_N)
 
-    topk_weights_data = tl.load(topk_weights + n_range, mask=n_range < experts_info_n, other=0)
-    expert_data = tl.load(
-        experts_info_ptr + expert_id * experts_info_stride0 + n_range, mask=n_range < experts_info_n, other=0
-    )
-    cumsum_expert_data = tl.cumsum(expert_data)
+    off_n = tl.arange(0, TOKEN_BLOCK_SIZE)
 
-    tl.store(expert_token_num_ptr + expert_id, tl.max(cumsum_expert_data))
-    tl.store(
-        experts_info_ptr + expert_id * experts_info_stride0 + cumsum_expert_data - 1,
-        n_range,
-        mask=(expert_data == 1) & (n_range < experts_info_n),
-    )
-    tl.store(
-        experts_topk_weight + expert_id * experts_topk_weight_stride0 + cumsum_expert_data - 1,
-        topk_weights_data,
-        mask=(expert_data == 1) & (n_range < experts_info_n),
-    )
+    pre_sum = 0
+
+    for start_loc in tl.range(0, experts_info_n, TOKEN_BLOCK_SIZE, num_stages=NUM_STAGE):
+        n_range = start_loc + off_n
+        topk_weights_data = tl.load(topk_weights + n_range, mask=n_range < experts_info_n, other=0)
+        expert_data = tl.load(
+            experts_info_ptr + expert_id * experts_info_stride0 + n_range, mask=n_range < experts_info_n, other=0
+        )
+        cumsum_expert_data = tl.cumsum(expert_data) + pre_sum
+        pre_sum = tl.max(cumsum_expert_data)
+        tl.store(
+            experts_info_ptr + expert_id * experts_info_stride0 + cumsum_expert_data - 1,
+            n_range,
+            mask=(expert_data == 1) & (n_range < experts_info_n),
+        )
+        tl.store(
+            experts_topk_weight + expert_id * experts_topk_weight_stride0 + cumsum_expert_data - 1,
+            topk_weights_data,
+            mask=(expert_data == 1) & (n_range < experts_info_n),
+        )
+
+    tl.store(expert_token_num_ptr + expert_id, pre_sum)
+    return
 
 
 def moe_align1(
@@ -184,7 +192,11 @@ def moe_align1(
     assert token_num_mul_topk <= FFN_MOE_CHUNK_SIZE * topk_num, "need split to handle seq len too long"
     assert exports_token_num.shape[0] == expert_num
     assert topk_weights.is_contiguous()
-    TOKEN_BLOCK_N = triton.next_power_of_2(token_num_mul_topk)
+    if token_num_mul_topk <= 512:
+        TOKEN_BLOCK_SIZE = 256
+    else:
+        TOKEN_BLOCK_SIZE = 512 if token_num_mul_topk <= 4 * 1024 else 2048
+
     grid = (expert_num,)
     moe_align1_kernel[grid](
         experts_info,
@@ -197,7 +209,8 @@ def moe_align1(
         experts_weight_info,
         experts_weight_info.stride(0),
         experts_weight_info.stride(1),
-        TOKEN_BLOCK_N=TOKEN_BLOCK_N,
+        TOKEN_BLOCK_SIZE=TOKEN_BLOCK_SIZE,
+        NUM_STAGE=4,
         num_warps=8,
         num_stages=1,
     )

…ommon/fused_moe/test_groped_fused_moe.py …ommon/fused_moe/test_groped_fused_moe.pylightllm/common/fused_moe/test_groped_fused_moe.py renamed to unit_tests/common/fused_moe/test_groped_fused_moe.py lightllm/common/fused_moe/test_groped_fused_moe.py renamed to unit_tests/common/fused_moe/test_groped_fused_moe.py

@@ -1,17 +1,18 @@
 import torch
 import time
-from .grouped_fused_moe import moe_align, moe_align1, grouped_matmul
+import pytest
+from lightllm.common.fused_moe.grouped_fused_moe import moe_align, moe_align1, grouped_matmul
 from lightllm.utils.log_utils import init_logger
 
+logger = init_logger(__name__)
+
 seed = 42
 torch.manual_seed(seed)
 
 if torch.cuda.is_available():
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
 
-logger = init_logger(__name__)
-
 
 def test_moe_align():
     expert_num = 5
@@ -137,6 +138,4 @@ def test_grouped_matmul():
 
 
 if __name__ == "__main__":
-    test_moe_align()
-    test_moe_align1()
-    test_grouped_matmul()
+    pytest.main()

unit_tests/common/fused_moe/test_groped_fused_moe_speed.py

@@ -0,0 +1,51 @@
+import torch
+import time
+import pytest
+from lightllm.common.fused_moe.grouped_fused_moe import moe_align, moe_align1, grouped_matmul
+from lightllm.utils.log_utils import init_logger
+
+seed = 42
+torch.manual_seed(seed)
+
+if torch.cuda.is_available():
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+logger = init_logger(__name__)
+
+
+@pytest.mark.parametrize("token_num", [200, 256, 8 * 1024])
+def test_moe_align1(token_num):
+    expert_num = 160
+    topk_num = 6
+    print(token_num)
+
+    def get_one():
+        rnd_logics = torch.randn(token_num, expert_num, device="cuda")
+        topk_values, topk_ids = torch.topk(rnd_logics, topk_num, dim=1)
+
+        experts_info = torch.zeros((expert_num, token_num * topk_num), dtype=torch.int32, device="cuda")
+        experts_info.fill_(0)
+        moe_align(topk_ids, experts_info)
+
+        topk_weights = torch.randn((token_num, topk_num), dtype=torch.float32, device="cuda")
+        experts_token_num = torch.zeros((expert_num,), dtype=torch.int32, device="cuda")
+        experts_weights = torch.zeros(experts_info.shape, dtype=torch.float32, device="cuda")
+        return experts_info, topk_weights, experts_weights, experts_token_num
+
+    test_datas = [get_one() for _ in range(100)]
+
+    moe_align1(*test_datas[0], topk_num)
+
+    torch.cuda.synchronize()
+    start = time.time()
+
+    for i in range(60):
+        moe_align1(*test_datas[i + 1], topk_num)
+    torch.cuda.synchronize()
+
+    print(f"token_num: {token_num} cost time: {time.time() - start} s")
+
+
+if __name__ == "__main__":
+    pytest.main()