moe_align_block_size op debug (#1240)

* fix(moe): Initialize output tensors and correct padding logic in moe_align_block_size_triton.

* add moe_align_block_size benchmark & test

* fix format

* fix format

* add comment

Author: huangyiqun

benchmark/test_special_perf.py

@@ -561,3 +561,72 @@ def torch_op(input_tensor, output_tensor):
     )
     bench.set_gems(gems_op)
     bench.run()
+
+
+try:
+    import os
+
+    os.environ["VLLM_CONFIGURE_LOGGING"] = "0"
+    import vllm._custom_ops as vllm_ops
+
+    HAS_VLLM = True
+except ImportError:
+    HAS_VLLM = False
+
+
+@pytest.mark.moe_align_block_size
+@pytest.mark.skipif(not HAS_VLLM, reason="vllm not installed")
+def test_perf_moe_align_block_size():
+    def moe_align_block_size_input_fn(shape, dtype, device):
+        # ------------ parameters ------------
+        num_experts = shape[0]
+        block_size = shape[1]
+        dtype = torch.int32
+        topk_ids = torch.randint(0, num_experts, (3, 4), dtype=dtype, device=device)
+        max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+
+        # padded_num_experts in vllm._custom_ops.moe_align_block_size
+        # must be less than 1024
+        if max_num_tokens_padded >= 1024:
+            return
+
+        sorted_ids = torch.empty((max_num_tokens_padded,), dtype=dtype, device=device)
+        max_num_m_blocks = max_num_tokens_padded // block_size
+        expert_ids = torch.empty((max_num_m_blocks,), dtype=dtype, device=device)
+        num_tokens_post_pad = torch.empty(1, dtype=dtype, device=device)
+
+        yield (
+            topk_ids,
+            num_experts,
+            block_size,
+            sorted_ids,
+            expert_ids,
+            num_tokens_post_pad,
+        )
+
+    class MoeAlignBlockSizeBenchmark(GenericBenchmark2DOnly):
+        def set_more_shapes(self):
+            return [
+                (16, 8),
+                (16, 16),
+                (16, 32),
+                (32, 8),
+                (32, 16),
+                (32, 32),
+                (64, 8),
+                (64, 16),
+                (128, 8),
+            ]
+
+    gems_op = flag_gems.moe_align_block_size_triton
+    bench = MoeAlignBlockSizeBenchmark(
+        op_name="moe_align_block_size_triton",
+        input_fn=moe_align_block_size_input_fn,
+        torch_op=vllm_ops.moe_align_block_size,
+        dtypes=[
+            torch.int32,
+        ],
+    )
+
+    bench.set_gems(gems_op)
+    bench.run()

src/flag_gems/fused/moe_align_block_size.py

@@ -107,6 +107,11 @@ def moe_align_block_size_triton(
     num_tokens_post_pad: torch.Tensor,
 ) -> None:
     numel = topk_ids.numel()
+    # The tensor needs to be padded before calculating IDs,
+    # to prevent out-of-bounds address access.
+    sorted_token_ids.fill_(numel)
+    expert_ids.fill_(0)
+
     grid = (num_experts,)
     tokens_cnts = torch.zeros(
         (num_experts + 1, num_experts), dtype=torch.int32, device=topk_ids.device

tests/test_DSA/test_bin_topk.py

@@ -121,7 +121,7 @@ def debug_topk_results(actual, expected, inputs, test_name=""):
         print(f"  Actual indices: {sorted(actual_set)[:m]}...")  # Only show first 10
         print(f"  Expected indices: {sorted(expected_set)[:m]}...")
         print(
-            f"  Intersection: {len(intersection)}/{len(expected_set)} = {len(intersection)/len(expected_set):.4f}"
+            f"  Intersection: {len(intersection)}/{len(expected_set)} = {len(intersection) / len(expected_set):.4f}"
         )
 
         # Check quality of actually selected values

tests/test_special_ops.py

@@ -1389,3 +1389,58 @@ def test_moe_sum(shape, dtype):
     with flag_gems.use_gems():
         flag_gems.moe_sum(inp1, res_out)
     gems_assert_close(res_out, ref_out, dtype)
+
+
+try:
+    import vllm._custom_ops as vllm_ops
+
+    HAS_VLLM = True
+except ImportError:
+    HAS_VLLM = False
+
+
+# ref: https://github.com/vllm-project/vllm/blob/main/tests/kernels/moe/test_moe.py
+@pytest.mark.moe_align_block_size
+@pytest.mark.parametrize("num_experts", [32, 256, 512])
+@pytest.mark.parametrize("block_size", [8, 16, 32])
+@pytest.mark.skipif(not HAS_VLLM, reason="vllm not installed")
+def test_accuracy_moe_align_block_size(
+    num_experts,
+    block_size,
+):
+    # ------------ parameters ------------
+    dtype = torch.int32
+    topk_ids = torch.randint(0, num_experts, (3, 4), dtype=dtype, device="cuda")
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    sorted_ids = torch.empty((max_num_tokens_padded,), dtype=dtype, device="cuda")
+    max_num_m_blocks = max_num_tokens_padded // block_size
+    expert_ids = torch.empty((max_num_m_blocks,), dtype=dtype, device="cuda")
+    num_tokens_post_pad = torch.empty(1, dtype=dtype, device="cuda")
+
+    topk_ids_vllm = topk_ids.clone()
+    sorted_ids_vllm = sorted_ids.clone()
+    expert_ids_vllm = expert_ids.clone()
+    num_tokens_post_pad_vllm = num_tokens_post_pad.clone()
+
+    flag_gems.moe_align_block_size_triton(
+        topk_ids=topk_ids,
+        num_experts=num_experts,
+        block_size=block_size,
+        sorted_token_ids=sorted_ids,
+        expert_ids=expert_ids,
+        num_tokens_post_pad=num_tokens_post_pad,
+    )
+
+    vllm_ops.moe_align_block_size(
+        topk_ids=topk_ids_vllm,
+        num_experts=num_experts,
+        block_size=block_size,
+        sorted_token_ids=sorted_ids_vllm,
+        experts_ids=expert_ids_vllm,
+        num_tokens_post_pad=num_tokens_post_pad_vllm,
+    )
+
+    torch.cuda.synchronize()
+    gems_assert_close(sorted_ids, sorted_ids_vllm, dtype=dtype)
+    gems_assert_close(expert_ids, expert_ids_vllm, dtype=dtype)
+    gems_assert_close(num_tokens_post_pad, num_tokens_post_pad_vllm, dtype=dtype)