A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Bug] Fix moe_sum signature (#18440)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
May 21, 2025 at 01:37 UTC
92247c522e216f9d010db1c648dc783dbf141704
2 files changed
+19
-1
csrc/moe/torch_bindings.cpp
+1
-1
@@ -10,7 +10,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
10
11
// Calculate the result of moe by summing up the partial results
12
// from all selected experts.
13
- m.def("moe_sum(Tensor! input, Tensor output) -> ()");
13
+ m.def("moe_sum(Tensor input, Tensor! output) -> ()");
14
m.impl("moe_sum", torch::kCUDA, &moe_sum);
15
16
// Aligning the number of tokens to be processed by each expert such
tests/kernels/moe/test_moe.py
+18
@@ -575,3 +575,21 @@ def test_moe_align_block_size_opcheck():
575
opcheck(torch.ops._moe_C.moe_align_block_size,
576
(topk_ids, num_experts, block_size, sorted_ids, expert_ids,
577
num_tokens_post_pad))
578
+
579
+
580
+@pytest.mark.parametrize("m", [1, 33, 222, 1024 * 128])
581
+@pytest.mark.parametrize("topk", TOP_KS)
582
+@pytest.mark.parametrize("k", [128, 511, 1024])
583
+@pytest.mark.parametrize("dtype",
584
+ [torch.float32, torch.float16, torch.bfloat16])
585
+@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
586
+def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
587
+ input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
588
+ actual = torch.empty((m, k), device="cuda", dtype=dtype)
589
+
590
+ expected = input.sum(dim=1)
591
+ torch.ops._moe_C.moe_sum(input, actual)
592
+
593
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
594
+
595
+ opcheck(torch.ops._moe_C.moe_sum, (input, actual))