Improve torch compile for fused moe (#2327)

2024-12-03 01:58:25 -08:00
parent 83b340e371
commit 07ec07ad1f
6 changed files with 45 additions and 24 deletions
--- a/test/srt/test_torch_compile_moe.py
+++ b/test/srt/test_torch_compile_moe.py
@@ -14,7 +14,7 @@ from sglang.test.test_utils import (
 )


-class TestTorchCompile(unittest.TestCase):
+class TestTorchCompileMoe(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.model = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST
@@ -23,7 +23,7 @@ class TestTorchCompile(unittest.TestCase):
            cls.model,
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=["--enable-torch-compile", "--torch-compile-max-bs", "1"],
+            other_args=["--enable-torch-compile", "--torch-compile-max-bs", "8"],
        )

    @classmethod