ci: migrate MLA tests to test/registered/mla/ (#15798)

.github/workflows/pr-test.yml

@@ -433,7 +433,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        partition: [0, 1, 2]
+        partition: [0, 1, 2, 3]
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -454,7 +454,7 @@ jobs:
         timeout-minutes: 30
         run: |
           cd test/
-          python3 run_suite.py --hw cuda --suite stage-b-test-small-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 3
+          python3 run_suite.py --hw cuda --suite stage-b-test-small-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 4
 
   stage-b-test-large-1-gpu:
     needs: [check-changes, call-gate, stage-a-test-1, sgl-kernel-build-wheels]

test/registered/mla/test_flashmla.py

@@ -1,6 +1,6 @@
 """
 Usage:
-python3 test/srt/test_flashmla.py
+python3 test/registered/mla/test_flashmla.py
 """
 
 import unittest
@@ -10,6 +10,7 @@ import requests
 import torch
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST_MLA,
@@ -19,6 +20,9 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+# FlashMLA attention backend tests with MTP speculative decoding
+register_cuda_ci(est_time=230, suite="stage-b-test-small-1-gpu")
+
 
 class TestFlashMLAAttnBackend(unittest.TestCase):
     @classmethod

test/registered/mla/test_mla.py

@@ -2,6 +2,7 @@ import unittest
 from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_MLA_MODEL_NAME_FOR_TEST,
@@ -11,6 +12,14 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+# MLA attention test with MGSM evaluation
+register_cuda_ci(est_time=194, suite="stage-b-test-small-1-gpu")
+register_amd_ci(
+    est_time=242,
+    suite="stage-a-test-1",
+    disabled="see https://github.com/sgl-project/sglang/issues/13107",
+)
+
 
 class TestMLA(CustomTestCase):
     @classmethod

test/registered/mla/test_mla_deepseek_v3.py

@@ -5,6 +5,7 @@ from types import SimpleNamespace
 import requests
 
 from sglang.srt.utils import is_cuda, is_hip, kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -14,6 +15,14 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+# DeepSeek-V3 MLA tests with torch compile, FA3, and MTP speculative decoding
+register_cuda_ci(est_time=442, suite="stage-b-test-small-1-gpu")
+register_amd_ci(
+    est_time=221,
+    suite="stage-a-test-1",
+    disabled="see https://github.com/sgl-project/sglang/issues/12574",
+)
+
 
 class TestMLADeepseekV3(CustomTestCase):
     @classmethod

test/registered/mla/test_mla_flashinfer.py

@@ -5,6 +5,7 @@ import requests
 import torch
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -13,6 +14,9 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+# FlashInfer MLA backend tests with MTP speculative decoding
+register_cuda_ci(est_time=302, suite="stage-b-test-small-1-gpu")
+
 
 class TestFlashinferMLA(CustomTestCase):
     @classmethod

test/registered/mla/test_mla_fp8.py

@@ -2,6 +2,7 @@ import unittest
 from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_MLA_FP8_MODEL_NAME_FOR_TEST,
@@ -11,6 +12,9 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+# MLA FP8 KV cache test with MGSM evaluation
+register_cuda_ci(est_time=77, suite="stage-b-test-small-1-gpu")
+
 
 class TestMLA(CustomTestCase):
     @classmethod

test/registered/mla/test_mla_int8_deepseek_v3.py

@@ -5,6 +5,7 @@ import requests
 import torch
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -14,6 +15,9 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+# DeepSeek-V3 INT8 quantization tests (channel and block INT8)
+register_cuda_ci(est_time=300, suite="stage-b-test-small-1-gpu")
+
 
 class TestMLADeepseekV3ChannelInt8(CustomTestCase):
     @classmethod

test/srt/run_suite.py

@@ -58,7 +58,6 @@ suites = {
         TestFile("test_constrained_decoding.py", 111),
         TestFile("test_eval_fp8_accuracy.py", 250),
         TestFile("test_external_models.py", 30),
-        TestFile("test_flashmla.py", 230),
         TestFile("test_fp8_utils.py", 9),
         TestFile("rotary_embedding/test_mrope.py", 10),
         TestFile("test_fused_moe.py", 80),
@@ -71,11 +70,6 @@ suites = {
         TestFile("test_mamba_unittest.py", 9),
         TestFile("test_metrics.py", 32),
         TestFile("test_metrics_utils.py", 1),
-        TestFile("test_mla.py", 194),
-        TestFile("test_mla_deepseek_v3.py", 442),
-        TestFile("test_mla_flashinfer.py", 302),
-        TestFile("test_mla_fp8.py", 77),
-        TestFile("test_mla_int8_deepseek_v3.py", 300),
         TestFile("test_model_hooks.py", 6),
         TestFile("test_modelopt_loader.py", 11),
         TestFile("test_multi_tokenizer.py", 230),
@@ -247,8 +241,6 @@ suite_amd = {
         TestFile("test_jinja_template_utils.py", 1),
         TestFile("test_metrics.py", 32),
         TestFile("test_metrics_utils.py", 1),
-        # TestFile("test_mla.py", 242), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
-        # TestFile("test_mla_deepseek_v3.py", 221), # Temporarily disabled, see https://github.com/sgl-project/sglang/issues/12574
         TestFile("test_no_chunked_prefill.py", 108),
         TestFile("test_page_size.py", 60),
         TestFile("test_penalty.py", 180),
@@ -275,7 +267,6 @@ suite_amd = {
     ],
     "per-commit-amd-mi35x": [
         TestFile("test_gpt_oss_1gpu.py", 750),
-        TestFile("test_mla.py", 242),
     ],
     "per-commit-2-gpu-amd": [
         # TestFile("lora/test_lora_tp.py", 116), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107. Moved to test/registered/lora/