minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Khoa Pham
2026-04-06 00:11:14 -07:00
committed by GitHub
parent b311db2e49
commit 12272b6791
17 changed files with 1026 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/spec/test_ngram_speculative_decoding.py
+81
View File
@@ -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()
test/registered/unit/server_args/test_server_args.py
+58
View File
@@ -426,5 +426,63 @@ class TestHiCacheArgs(unittest.TestCase):
self.assertEqual(args.decode_attention_backend, "triton")
class TestNgramExternalSamArgs(CustomTestCase):
def _make_dummy_ngram_args(self, **overrides):
args = ServerArgs(model_path="dummy")
args.speculative_algorithm = "NGRAM"
args.speculative_num_draft_tokens = 12
args.device = "cuda"
for key, value in overrides.items():
setattr(args, key, value)
return args
def test_prepare_server_args_parses_external_sam_args(self):
server_args = prepare_server_args(
[
"--model-path",
"dummy",
"--speculative-algorithm",
"NGRAM",
"--speculative-ngram-external-corpus-path",
"/tmp/ngram-corpus.jsonl",
"--speculative-ngram-external-sam-budget",
"4",
"--speculative-ngram-external-corpus-max-tokens",
"128",
]
)
self.assertEqual(
server_args.speculative_ngram_external_corpus_path,
"/tmp/ngram-corpus.jsonl",
)
self.assertEqual(server_args.speculative_ngram_external_sam_budget, 4)
self.assertEqual(server_args.speculative_ngram_external_corpus_max_tokens, 128)
def test_external_sam_budget_requires_path(self):
with self.assertRaises(ValueError) as context:
self._make_dummy_ngram_args(
speculative_ngram_external_sam_budget=2,
)._handle_speculative_decoding()
self.assertIn("external-sam-budget", str(context.exception))
def test_external_sam_budget_must_fit_draft_budget(self):
with self.assertRaises(ValueError) as context:
self._make_dummy_ngram_args(
speculative_num_draft_tokens=4,
speculative_ngram_external_corpus_path="/tmp/ngram-corpus.jsonl",
speculative_ngram_external_sam_budget=4,
)._handle_speculative_decoding()
self.assertIn("speculative_num_draft_tokens - 1", str(context.exception))
def test_external_corpus_max_tokens_must_be_positive(self):
with self.assertRaises(ValueError) as context:
self._make_dummy_ngram_args(
speculative_ngram_external_corpus_path="/tmp/ngram-corpus.jsonl",
speculative_ngram_external_sam_budget=2,
speculative_ngram_external_corpus_max_tokens=0,
)._handle_speculative_decoding()
self.assertIn("external-corpus-max-tokens", str(context.exception))
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_ngram_corpus.py
+168
View File
@@ -1,8 +1,14 @@
import json
import os
import tempfile
import unittest
import uuid
import numpy as np
from sglang.srt.speculative.cpp_ngram.external_corpus import (
iter_external_corpus_chunks,
)
from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -17,6 +23,9 @@ def _make_corpus(match_type="BFS", **kwargs):
max_bfs_breadth=8,
draft_token_num=8,
capacity=100000,
external_sam_budget=0,
external_corpus_max_tokens=10000000,
external_corpus_documents=None,
)
defaults.update(kwargs)
defaults["match_type"] = match_type
@@ -43,6 +52,12 @@ def _batch_get_with_state(
return corpus.batch_get([req_id], [current_tokens], [total_len])
class _IntTokenizer:
def encode(self, text: str, add_special_tokens: bool = False):
del add_special_tokens
return [int(piece) for piece in text.split()]
SEED_SEQUENCES = [
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[1, 2, 3, 44, 55, 66, 77, 88, 99, 100],
@@ -674,6 +689,159 @@ class TestNgramCorpusIncremental(CustomTestCase):
np.testing.assert_array_equal(inc_masks, full_masks)
class TestNgramCorpusExternalSam(CustomTestCase):
"""Verify external SAM loading and fixed-budget composition."""
def test_external_corpus_iterator_streams_documents(self):
corpus = _make_corpus(
"BFS",
draft_token_num=4,
external_sam_budget=3,
external_corpus_max_tokens=8,
)
with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
f.write(json.dumps("1 2 3 4 5"))
f.write("\n")
f.write(json.dumps("8 9"))
f.write("\n")
path = f.name
self.addCleanup(os.remove, path)
loaded_token_count = corpus.load_external_corpus(
iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=8)
)
# 5 doc tokens + 1 separator + 2 doc tokens = 8
self.assertEqual(loaded_token_count, 8)
ids, _ = _batch_get(corpus, [[1, 2, 3]])
ids_list = ids.tolist()
self.assertEqual(ids_list[0], 3)
self.assertEqual(ids_list[1:3], [4, 5])
def test_external_corpus_iterator_rejects_oversized_corpus(self):
corpus = _make_corpus(
"BFS",
external_sam_budget=2,
external_corpus_max_tokens=4,
)
with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f:
f.write(json.dumps("1 2 3"))
f.write("\n")
f.write(json.dumps("4 5"))
f.write("\n")
path = f.name
self.addCleanup(os.remove, path)
with self.assertRaisesRegex(ValueError, "token limit"):
corpus.load_external_corpus(
iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=4)
)
def test_external_sam_documents_reject_oversized_corpus(self):
with self.assertRaisesRegex(ValueError, "token limit"):
_make_corpus(
"BFS",
external_sam_budget=2,
external_corpus_max_tokens=4,
external_corpus_documents=[[1, 2, 3], [4, 5]],
)
def test_external_sam_only_chain(self):
corpus = _make_corpus(
"BFS",
draft_token_num=4,
external_sam_budget=3,
external_corpus_documents=[[1, 2, 3, 4, 5]],
)
ids, masks = _batch_get(corpus, [[1, 2, 3]])
ids_list = ids.tolist()
self.assertEqual(ids_list[0], 3)
self.assertEqual(ids_list[1:3], [4, 5])
def test_external_sam_respects_document_boundaries(self):
corpus = _make_corpus(
"BFS",
draft_token_num=4,
external_sam_budget=3,
external_corpus_documents=[[1, 2, 3], [4, 5, 6]],
)
ids, _ = _batch_get(corpus, [[2, 3]])
ids_list = ids.tolist()
self.assertEqual(ids_list[0], 3)
self.assertTrue(all(token == 0 for token in ids_list[1:]), ids_list)
def test_external_sam_adds_distinct_root_branch(self):
corpus = _make_corpus(
"BFS",
draft_token_num=6,
external_sam_budget=2,
external_corpus_documents=[[1, 2, 3, 20, 21]],
)
corpus.batch_put([[1, 2, 3, 10, 11]])
corpus.synchronize()
ids, masks = _batch_get(corpus, [[1, 2, 3]])
leaf_paths = corpus.leaf_paths_from_mask(
ids.tolist(), masks.reshape(6, 6).tolist()
)
self.assertIn([3, 10, 11], leaf_paths)
self.assertIn([3, 20, 21], leaf_paths)
def test_shared_prefix_keeps_both_branches(self):
corpus = _make_corpus(
"BFS",
draft_token_num=5,
external_sam_budget=2,
external_corpus_documents=[[1, 2, 3, 10, 99]],
)
corpus.batch_put([[1, 2, 3, 10, 11]])
corpus.synchronize()
ids, masks = _batch_get(corpus, [[1, 2, 3]])
leaf_paths = corpus.leaf_paths_from_mask(
ids.tolist(), masks.reshape(5, 5).tolist()
)
self.assertIn([3, 10, 11], leaf_paths)
self.assertIn([3, 10, 99], leaf_paths)
def test_shared_prefix_merge_can_underfill_budget(self):
corpus = _make_corpus(
"BFS",
draft_token_num=6,
external_sam_budget=2,
external_corpus_documents=[[1, 2, 3, 10, 99]],
)
corpus.batch_put([[1, 2, 3, 10, 11]])
corpus.synchronize()
ids, masks = _batch_get(corpus, [[1, 2, 3]])
ids_list = ids.tolist()
leaf_paths = corpus.leaf_paths_from_mask(ids_list, masks.reshape(6, 6).tolist())
self.assertIn([3, 10, 11], leaf_paths)
self.assertIn([3, 10, 99], leaf_paths)
self.assertEqual(ids_list.count(0), 2, ids_list)
def test_external_sam_prob_prefers_frequent_continuation(self):
corpus = _make_corpus(
"PROB",
draft_token_num=2,
min_bfs_breadth=1,
max_bfs_breadth=1,
external_sam_budget=1,
external_corpus_documents=[
[1, 2, 3, 10],
[1, 2, 3, 20],
[1, 2, 3, 20],
[1, 2, 3, 20],
],
)
ids, _ = _batch_get(corpus, [[1, 2, 3]])
self.assertEqual(ids.tolist(), [3, 20])
class TestNgramCorpusMatchBenchmark(CustomTestCase):
"""Benchmark incremental advance vs full rebuild in match()."""