[Spec][Ngram] 6/N: Load an external corpus and construct a Suffix Automaton (#21425)

test/registered/spec/test_ngram_speculative_decoding.py

@@ -1,5 +1,10 @@
+import json
+import os
+import tempfile
 import unittest
 
+import requests
+
 from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -30,6 +35,21 @@ DEFAULT_SERVER_ARGS = [
     0.8,
 ]
 
+EXTERNAL_SAM_CORPUS_RECORDS = [
+    "The capital of France is Paris.",
+    "The answer to life, the universe, and everything is 42.",
+]
+
+
+def _safe_remove(path: str):
+    if os.path.exists(path):
+        os.remove(path)
+
+
+def _safe_kill_process(process):
+    if process is not None and process.poll() is None:
+        kill_process_tree(process.pid)
+
 
 class TestNgramSpeculativeDecodingBase(GSM8KMixin, CustomTestCase):
     model = DEFAULT_TARGET_MODEL_NGRAM
@@ -86,5 +106,66 @@ class TestNgramSpeculativeDecodingPaged(TestNgramSpeculativeDecodingBase):
         ]
 
 
+class TestNgramExternalSamSmoke(CustomTestCase):
+    model = DEFAULT_TARGET_MODEL_NGRAM
+    base_url = DEFAULT_URL_FOR_TEST
+    attention_backends = ("triton", "flashinfer")
+
+    def get_server_args(self, attention_backend):
+        return DEFAULT_SERVER_ARGS + [
+            "--attention-backend",
+            attention_backend,
+            "--speculative-ngram-external-corpus-path",
+            self.external_corpus_path,
+            "--speculative-ngram-external-sam-budget",
+            "4",
+        ]
+
+    @classmethod
+    def setUpClass(cls):
+        envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.set(False)
+        envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".jsonl", prefix="ngram_external_sam_", delete=False
+        ) as f:
+            for record in EXTERNAL_SAM_CORPUS_RECORDS:
+                f.write(json.dumps(record))
+                f.write("\n")
+            cls.external_corpus_path = f.name
+        cls.addClassCleanup(_safe_remove, cls.external_corpus_path)
+
+    def _run_external_sam_smoke(self, attention_backend):
+        process = popen_launch_server(
+            self.model,
+            self.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=self.get_server_args(attention_backend),
+        )
+        try:
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": "The capital of France is",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": 8,
+                    },
+                },
+                timeout=120,
+            )
+            self.assertEqual(response.status_code, 200, response.text)
+            response_json = response.json()
+            self.assertIn("text", response_json)
+            self.assertIn("meta_info", response_json)
+            self.assertGreater(response_json["meta_info"]["completion_tokens"], 0)
+        finally:
+            _safe_kill_process(process)
+
+    def test_generate_with_external_sam(self):
+        for attention_backend in self.attention_backends:
+            with self.subTest(attention_backend=attention_backend):
+                self._run_external_sam_smoke(attention_backend)
+
+
 if __name__ == "__main__":
     unittest.main()