Import flash_mla from sgl-kernel (#12135)

test/srt/run_suite.py

@@ -82,7 +82,7 @@ suites = {
         TestFile("test_ebnf_constrained.py", 108),
         TestFile("test_eval_fp8_accuracy.py", 303),
         TestFile("test_fa3.py", 376),
-        # TestFile("test_flashmla.py", 352),
+        TestFile("test_flashmla.py", 352),
         TestFile("rotary_embedding/test_mrope.py", 300),
         TestFile("test_function_call_parser.py", 10),
         TestFile("test_fused_moe.py", 30),

test/srt/test_flashmla.py

@@ -11,6 +11,7 @@ import torch
 
 from sglang.srt.utils import kill_process_tree
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.send_one import BenchArgs, send_one_prompt
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST_MLA,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -18,7 +19,7 @@ from sglang.test.test_utils import (
     CustomTestCase,
     is_in_ci,
     popen_launch_server,
-    run_bench_one_batch,
+    write_github_step_summary,
 )
 
 
@@ -31,7 +32,6 @@ class TestFlashMLAAttnBackend(unittest.TestCase):
         if torch.cuda.is_available() and torch.version.cuda:
             other_args.extend(
                 [
-                    "--enable-torch-compile",
                     "--cuda-graph-max-bs",
                     "2",
                     "--attention-backend",
@@ -65,24 +65,6 @@ class TestFlashMLAAttnBackend(unittest.TestCase):
         self.assertGreater(metrics["accuracy"], 0.60)
 
 
-class TestFlashMLAAttnLatency(unittest.TestCase):
-    def test_latency(self):
-        _, output_throughput, _ = run_bench_one_batch(
-            DEFAULT_MODEL_NAME_FOR_TEST_MLA,
-            [
-                "--attention-backend",
-                "flashmla",
-                "--enable-torch-compile",
-                "--cuda-graph-max-bs",
-                "16",
-                "--trust-remote-code",
-            ],
-        )
-
-        if is_in_ci():
-            self.assertGreater(output_throughput, 100)
-
-
 class TestFlashMLAMTP(CustomTestCase):
     @classmethod
     def setUpClass(cls):