[Kimi K2] dsv3_router_gemm supports NUM_EXPERTS == 384 (#8013)

2025-08-01 22:01:24 +08:00
parent 46e9d1c7c1
commit 6bdd27861b
5 changed files with 188 additions and 30 deletions
--- a/sgl-kernel/tests/test_dsv3_router_gemm.py
+++ b/sgl-kernel/tests/test_dsv3_router_gemm.py
@@ -5,8 +5,8 @@ from sgl_kernel import dsv3_router_gemm


@pytest.mark.parametrize("num_tokens", [i + 1 for i in range(16)])
-def test_dsv3_router_gemm(num_tokens):
-    num_experts = 256
+@pytest.mark.parametrize("num_experts", [256, 384])
+def test_dsv3_router_gemm(num_tokens, num_experts):
    hidden_dim = 7168

    mat_a = torch.randn(