Skip to content

Commit 1bb7259

Browse files
authored
add test case for trtllm gen fused moe with kimi k2 problem sizes (#1768)
This is to add kimi k2 problem size to unit test. num of experts is 384 top_k is 8 The support for kimi k2 moe is added w/ [PR 1696](#1696) ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. -->
1 parent b7be894 commit 1bb7259

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

tests/test_trtllm_gen_fused_moe.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1853,6 +1853,23 @@ def cache_permute_indices():
18531853
@pytest.mark.parametrize(
18541854
"routing_config",
18551855
[
1856+
pytest.param(
1857+
{
1858+
"num_experts": 384,
1859+
"top_k": 8,
1860+
"padding": 8,
1861+
"n_groups": 12,
1862+
"top_k_groups": 4,
1863+
"routed_scaling": 2.5,
1864+
"has_routing_bias": True,
1865+
"routing_method_type": RoutingMethodType.DeepSeekV3,
1866+
"compatible_moe_impls": [
1867+
FP4Moe,
1868+
FP8BlockScaleMoe,
1869+
],
1870+
},
1871+
id="kimi_k2",
1872+
),
18561873
pytest.param(
18571874
{
18581875
"num_experts": 256,

0 commit comments

Comments
 (0)