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

[Spec][Ngram] Support multiple SAMs with dynamic HTTP API (#22203)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-06 18:49:22 -07:00
committed by GitHub
parent 49cb7d546e
commit e4b1366a46
14 changed files with 685 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp
+59 -20
View File
@@ -63,27 +63,57 @@ void Ngram::asyncInsert(std::vector<std::vector<int32_t>>&& tokens) {
}
}
// NOTE: staging operations (start/append/finish) are called from a background
// thread during async corpus loading. They do NOT hold mutex_ because
// staging_sam_ is disjoint from sams_ / trie_. Only finishExternalCorpusLoad
// briefly acquires mutex_ when moving the completed SAM into sams_.
void Ngram::startExternalCorpusLoad() {
std::unique_lock<std::mutex> lock(mutex_);
sam_ = std::make_unique<SuffixAutomaton>();
if (staging_sam_) {
throw std::runtime_error("startExternalCorpusLoad called while another load is in progress");
}
staging_sam_ = std::make_unique<SuffixAutomaton>();
}
void Ngram::appendExternalCorpusTokens(const std::vector<int32_t>& tokens) {
std::unique_lock<std::mutex> lock(mutex_);
sam_->appendTokens(tokens);
if (!staging_sam_) {
throw std::runtime_error("appendExternalCorpusTokens called without startExternalCorpusLoad");
}
staging_sam_->appendTokens(tokens);
}
void Ngram::finishExternalCorpusLoad() {
std::unique_lock<std::mutex> lock(mutex_);
sam_->finalize();
if (sam_->empty()) {
sam_.reset();
void Ngram::finishExternalCorpusLoad(const std::string& corpus_id) {
if (!staging_sam_) {
throw std::runtime_error("finishExternalCorpusLoad called without startExternalCorpusLoad");
}
staging_sam_->finalize();
if (staging_sam_->empty()) {
staging_sam_.reset();
throw std::runtime_error("External corpus is empty — no tokens were loaded.");
}
// Only lock briefly to install the completed SAM.
std::unique_lock<std::mutex> lock(mutex_);
sams_[corpus_id] = std::move(staging_sam_);
}
void Ngram::removeExternalCorpus(const std::string& corpus_id) {
std::unique_lock<std::mutex> lock(mutex_);
sams_.erase(corpus_id);
}
void Ngram::clearExternalCorpus() {
std::unique_lock<std::mutex> lock(mutex_);
sam_.reset();
sams_.clear();
staging_sam_.reset();
}
std::vector<std::string> Ngram::listExternalCorpora() const {
std::unique_lock<std::mutex> lock(mutex_);
std::vector<std::string> ids;
ids.reserve(sams_.size());
for (const auto& [id, _] : sams_) {
ids.push_back(id);
}
return ids;
}
void Ngram::insertWorker() {
@@ -154,6 +184,14 @@ Result Ngram::batchMatch(
throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'.");
}
// All budget values are loop-invariant (mutex_ held, sams_ won't change).
const size_t num_sams = sams_.size();
const auto total_draft_token_num = param_.get_draft_token_num(tokens.size());
const size_t total_sam_budget =
num_sams > 0 ? std::min(param_.external_sam_budget, total_draft_token_num) : size_t{0};
const size_t per_sam_budget = num_sams > 0 ? total_sam_budget / num_sams : size_t{0};
const size_t trie_budget = total_draft_token_num - (per_sam_budget * num_sams);
Result merged;
for (size_t i = 0; i < state_ids.size(); ++i) {
const auto& suffix = tokens[i];
@@ -162,12 +200,8 @@ Result Ngram::batchMatch(
}
auto& state = match_state_[state_ids[i]];
const auto total_draft_token_num = param_.get_draft_token_num(tokens.size());
const auto sam_budget =
sam_ && !sam_->empty() ? std::min(param_.external_sam_budget, total_draft_token_num) : size_t{0};
const auto trie_budget = total_draft_token_num - sam_budget;
if (sam_budget == 0) {
if (total_sam_budget == 0 || per_sam_budget == 0) {
auto res = (trie_.get()->*trie_result_build_fn)(
suffix.data(), suffix.size(), suffix.back(), total_draft_token_num, param_, state, total_lens[i]);
merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
@@ -175,12 +209,17 @@ Result Ngram::batchMatch(
continue;
}
auto trie_res = (trie_.get()->*trie_result_build_fn)(
auto combined = (trie_.get()->*trie_result_build_fn)(
suffix.data(), suffix.size(), suffix.back(), trie_budget, param_, state, total_lens[i]);
auto sam_res = (sam_.get()->*sam_result_build_fn)(suffix.data(), suffix.size(), suffix.back(), sam_budget, param_);
auto res = combineRootResults_(suffix.back(), static_cast<int>(total_draft_token_num + 1), trie_res, sam_res);
merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());
for (const auto& [_, sam] : sams_) {
auto sam_res =
(sam.get()->*sam_result_build_fn)(suffix.data(), suffix.size(), suffix.back(), per_sam_budget, param_);
combined = combineRootResults_(suffix.back(), static_cast<int>(total_draft_token_num + 1), combined, sam_res);
}
merged.token.insert(merged.token.end(), combined.token.begin(), combined.token.end());
merged.mask.insert(merged.mask.end(), combined.mask.begin(), combined.mask.end());
}
return merged;
}
python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h
+13 -5
View File
@@ -19,12 +19,16 @@ namespace ngram {
class Ngram {
std::unique_ptr<Trie> trie_;
std::unique_ptr<SuffixAutomaton> sam_;
std::unordered_map<std::string, std::unique_ptr<SuffixAutomaton>> sams_;
// FIXME: single staging slot — only one corpus can be loaded at a time.
// To support concurrent loads, move staging into a per-load local variable.
std::unique_ptr<SuffixAutomaton> staging_sam_;
Param param_;
// NOTE: protects trie_ and pending_count_. Ensures batchMatch never reads
// trie_ while insertWorker is writing. After synchronize(), no pending
// inserts remain so mutex_ contention is effectively zero.
// NOTE: protects trie_, sams_, and pending_count_. staging_sam_ is NOT
// protected by mutex_ — it is only accessed from the corpus loading thread.
// finishExternalCorpusLoad briefly acquires mutex_ to move the completed
// SAM into sams_.
mutable std::mutex mutex_;
mutable std::condition_variable sync_cv_;
// NOTE: tracks inserts from enqueue through trie_->insert() completion,
@@ -46,10 +50,14 @@ class Ngram {
void appendExternalCorpusTokens(const std::vector<int32_t>& tokens);
void finishExternalCorpusLoad();
void finishExternalCorpusLoad(const std::string& corpus_id);
void removeExternalCorpus(const std::string& corpus_id);
void clearExternalCorpus();
std::vector<std::string> listExternalCorpora() const;
Result batchMatch(const std::vector<std::vector<int32_t>>& tokens);
Result batchMatch(
python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp
+18 -2
View File
@@ -112,14 +112,28 @@ struct NgramCorpusObj : public tvm::ffi::Object {
ngram_->appendExternalCorpusTokens(tokens);
}
void finish_external_corpus_load() {
ngram_->finishExternalCorpusLoad();
void finish_external_corpus_load(const std::string& corpus_id) {
ngram_->finishExternalCorpusLoad(corpus_id);
}
void remove_external_corpus(const std::string& corpus_id) {
ngram_->removeExternalCorpus(corpus_id);
}
void clear_external_corpus() {
ngram_->clearExternalCorpus();
}
std::string list_external_corpora() {
auto ids = ngram_->listExternalCorpora();
std::string result;
for (size_t i = 0; i < ids.size(); ++i) {
if (i > 0) result += "\n";
result += ids[i];
}
return result;
}
void synchronize() {
ngram_->synchronize();
}
@@ -161,7 +175,9 @@ void register_ngram_corpus() {
.def("start_external_corpus_load", &NgramCorpusObj::start_external_corpus_load)
.def("append_external_corpus_tokens", &NgramCorpusObj::append_external_corpus_tokens)
.def("finish_external_corpus_load", &NgramCorpusObj::finish_external_corpus_load)
.def("remove_external_corpus", &NgramCorpusObj::remove_external_corpus)
.def("clear_external_corpus", &NgramCorpusObj::clear_external_corpus)
.def("list_external_corpora", &NgramCorpusObj::list_external_corpora)
.def("synchronize", &NgramCorpusObj::synchronize)
.def("reset", &NgramCorpusObj::reset);
}
python/sglang/jit_kernel/ngram_corpus.py
+12 -3
View File
@@ -118,8 +118,8 @@ def get_ngram_corpus_cls():
state_ids_t = torch.tensor(state_ids, dtype=torch.int64)
self.erase_match_state(state_ids_t) # type: ignore
def load_external_corpus(
self, chunks: Iterable[Sequence[int]]
def load_external_corpus_named(
self, corpus_id: str, chunks: Iterable[Sequence[int]]
) -> Tuple[int, int]:
self.start_external_corpus_load() # type: ignore
chunk_count = 0
@@ -130,10 +130,19 @@ def get_ngram_corpus_cls():
loaded_token_count += len(tokens_t)
self.append_external_corpus_tokens(tokens_t) # type: ignore
chunk_count += 1
self.finish_external_corpus_load() # type: ignore
self.finish_external_corpus_load(corpus_id) # type: ignore
except Exception:
self.clear_external_corpus() # type: ignore
raise
return chunk_count, loaded_token_count
def remove_corpus(self, corpus_id: str) -> None:
self.remove_external_corpus(corpus_id) # type: ignore
def list_corpora(self) -> List[str]:
result = self.list_external_corpora() # type: ignore
if not result:
return []
return result.split("\n")
return NgramCorpusFFI
python/sglang/srt/entrypoints/http_server.py
+58
View File
@@ -733,6 +733,64 @@ async def flush_cache(timeout: float = Query(0.0, ge=0.0)):
)
@app.post("/add_external_corpus")
@auth_level(AuthLevel.ADMIN_OPTIONAL)
async def add_external_corpus(request: Request):
"""Add an external corpus for ngram speculative decoding."""
from sglang.srt.managers.io_struct import AddExternalCorpusReqInput
try:
obj = AddExternalCorpusReqInput(**(await request.json()))
except TypeError as e:
return ORJSONResponse(
{"success": False, "message": str(e)},
status_code=HTTPStatus.BAD_REQUEST,
)
result = await _global_state.tokenizer_manager.add_external_corpus(obj)
return ORJSONResponse(
{
"success": result.success,
"corpus_id": result.corpus_id,
"message": result.message,
"loaded_token_count": result.loaded_token_count,
},
status_code=200 if result.success else HTTPStatus.BAD_REQUEST,
)
@app.post("/remove_external_corpus")
@auth_level(AuthLevel.ADMIN_OPTIONAL)
async def remove_external_corpus(request: Request):
"""Remove an external corpus by ID."""
body = await request.json()
corpus_id = body.get("corpus_id")
if not corpus_id:
return ORJSONResponse(
{"success": False, "message": "corpus_id is required."},
status_code=HTTPStatus.BAD_REQUEST,
)
result = await _global_state.tokenizer_manager.remove_external_corpus(corpus_id)
return ORJSONResponse(
{"success": result.success, "message": result.message},
status_code=200 if result.success else HTTPStatus.BAD_REQUEST,
)
@app.get("/list_external_corpora")
@auth_level(AuthLevel.ADMIN_OPTIONAL)
async def list_external_corpora():
"""List all active external corpora."""
result = await _global_state.tokenizer_manager.list_external_corpora()
return ORJSONResponse(
{
"success": result.success,
"corpus_ids": result.corpus_ids,
"message": result.message,
},
status_code=200 if result.success else HTTPStatus.BAD_REQUEST,
)
@app.api_route("/clear_hicache_storage_backend", methods=["GET", "POST"])
@auth_level(AuthLevel.ADMIN_OPTIONAL)
async def clear_hicache_storage_backend_deprecated():
python/sglang/srt/managers/io_struct.py
+39
View File
@@ -1131,6 +1131,45 @@ class FlushCacheReqOutput(BaseReq):
message: str = ""
@dataclass
class AddExternalCorpusReqInput(BaseReq):
corpus_id: Optional[str] = None
file_path: Optional[str] = None
documents: Optional[List[str]] = None
token_chunks: Optional[List[List[int]]] = None
@dataclass
class AddExternalCorpusReqOutput(BaseReq):
success: bool
corpus_id: str = ""
message: str = ""
loaded_token_count: int = 0
@dataclass
class RemoveExternalCorpusReqInput(BaseReq):
corpus_id: str
@dataclass
class RemoveExternalCorpusReqOutput(BaseReq):
success: bool
message: str = ""
@dataclass
class ListExternalCorporaReqInput(BaseReq):
pass
@dataclass
class ListExternalCorporaReqOutput(BaseReq):
success: bool
corpus_ids: List[str] = field(default_factory=list)
message: str = ""
@dataclass
class AttachHiCacheStorageReqInput(BaseReq):
"""Dynamically attach (enable) HiCache storage backend at runtime.
python/sglang/srt/managers/scheduler.py
+60
View File
@@ -83,6 +83,8 @@ from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
from sglang.srt.managers.io_struct import (
AbortReq,
ActiveRanksOutput,
AddExternalCorpusReqInput,
AddExternalCorpusReqOutput,
AttachHiCacheStorageReqInput,
AttachHiCacheStorageReqOutput,
BaseBatchReq,
@@ -114,6 +116,8 @@ from sglang.srt.managers.io_struct import (
InitWeightsSendGroupForRemoteInstanceReqInput,
InitWeightsSendGroupForRemoteInstanceReqOutput,
InitWeightsUpdateGroupReqInput,
ListExternalCorporaReqInput,
ListExternalCorporaReqOutput,
LoadLoRAAdapterFromTensorsReqInput,
LoadLoRAAdapterFromTensorsReqOutput,
LoadLoRAAdapterReqInput,
@@ -122,6 +126,8 @@ from sglang.srt.managers.io_struct import (
PauseGenerationReqInput,
ProfileReq,
ReleaseMemoryOccupationReqInput,
RemoveExternalCorpusReqInput,
RemoveExternalCorpusReqOutput,
ResumeMemoryOccupationReqInput,
RpcReqInput,
RpcReqOutput,
@@ -599,6 +605,7 @@ class Scheduler(
def maybe_init_draft_worker(self):
if self.spec_algorithm.is_none():
self.draft_worker = None
self.external_corpus_manager = None
return
# Launch a draft worker for speculative decoding
@@ -625,6 +632,18 @@ class Scheduler(
DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args)
self.draft_worker = DraftWorkerClass(**draft_worker_kwargs)
if self.spec_algorithm.is_ngram():
from sglang.srt.speculative.external_corpus_manager import (
ExternalCorpusManager,
)
self.external_corpus_manager = ExternalCorpusManager(
self.draft_worker,
self.send_to_tokenizer.send_output,
)
else:
self.external_corpus_manager = None
def init_model_worker(self):
self.init_tp_model_worker()
self.maybe_init_draft_worker()
@@ -1251,6 +1270,15 @@ class Scheduler(
(PauseGenerationReqInput, self.pause_generation),
(ContinueGenerationReqInput, self.continue_generation),
(DumperControlReqInput, self.handle_dumper_control),
(AddExternalCorpusReqInput, self.add_external_corpus),
(
RemoveExternalCorpusReqInput,
self.remove_external_corpus,
),
(
ListExternalCorporaReqInput,
self.list_external_corpora,
),
]
)
@@ -1605,6 +1633,8 @@ class Scheduler(
self.recv_from_rpc.send_pyobj(output)
self._check_pending_flush()
if self.external_corpus_manager is not None:
self.external_corpus_manager.check_pending_load()
def init_req_max_new_tokens(self, req):
req.sampling_params.max_new_tokens = min(
@@ -2865,6 +2895,36 @@ class Scheduler(
pending_req,
)
def add_external_corpus(
self, recv_req: AddExternalCorpusReqInput
) -> Optional[AddExternalCorpusReqOutput]:
if self.external_corpus_manager is None:
return AddExternalCorpusReqOutput(
success=False,
message="Ngram speculative decoding is not enabled.",
)
return self.external_corpus_manager.add(recv_req)
def remove_external_corpus(
self, recv_req: RemoveExternalCorpusReqInput
) -> RemoveExternalCorpusReqOutput:
if self.external_corpus_manager is None:
return RemoveExternalCorpusReqOutput(
success=False,
message="Ngram speculative decoding is not enabled.",
)
return self.external_corpus_manager.remove(recv_req)
def list_external_corpora(
self, recv_req: ListExternalCorporaReqInput
) -> ListExternalCorporaReqOutput:
if self.external_corpus_manager is None:
return ListExternalCorporaReqOutput(
success=False,
message="Ngram speculative decoding is not enabled.",
)
return self.external_corpus_manager.list(recv_req)
def flush_cache_wrapped(
self, recv_req: FlushCacheReqInput
) -> Optional[FlushCacheReqOutput]:
python/sglang/srt/managers/tokenizer_communicator_mixin.py
+111
View File
@@ -23,6 +23,8 @@ import fastapi
import zmq
from sglang.srt.managers.io_struct import (
AddExternalCorpusReqInput,
AddExternalCorpusReqOutput,
AttachHiCacheStorageReqInput,
AttachHiCacheStorageReqOutput,
CheckWeightsReqInput,
@@ -53,6 +55,8 @@ from sglang.srt.managers.io_struct import (
InitWeightsSendGroupForRemoteInstanceReqOutput,
InitWeightsUpdateGroupReqInput,
InitWeightsUpdateGroupReqOutput,
ListExternalCorporaReqInput,
ListExternalCorporaReqOutput,
LoadLoRAAdapterFromTensorsReqInput,
LoadLoRAAdapterFromTensorsReqOutput,
LoadLoRAAdapterReqInput,
@@ -64,6 +68,8 @@ from sglang.srt.managers.io_struct import (
ProfileReqType,
ReleaseMemoryOccupationReqInput,
ReleaseMemoryOccupationReqOutput,
RemoveExternalCorpusReqInput,
RemoveExternalCorpusReqOutput,
ResumeMemoryOccupationReqInput,
ResumeMemoryOccupationReqOutput,
SendWeightsToRemoteInstanceReqInput,
@@ -205,6 +211,15 @@ class TokenizerCommunicatorMixin:
self.flush_cache_communicator = _Communicator(
self.send_to_scheduler, server_args.dp_size
)
self.add_external_corpus_communicator = _Communicator(
self.send_to_scheduler, server_args.dp_size
)
self.remove_external_corpus_communicator = _Communicator(
self.send_to_scheduler, server_args.dp_size
)
self.list_external_corpora_communicator = _Communicator(
self.send_to_scheduler, server_args.dp_size
)
self.clear_hicache_storage_communicator = _Communicator(
self.send_to_scheduler, server_args.dp_size
)
@@ -308,6 +323,18 @@ class TokenizerCommunicatorMixin:
FlushCacheReqOutput,
self.flush_cache_communicator.handle_recv,
),
(
AddExternalCorpusReqOutput,
self.add_external_corpus_communicator.handle_recv,
),
(
RemoveExternalCorpusReqOutput,
self.remove_external_corpus_communicator.handle_recv,
),
(
ListExternalCorporaReqOutput,
self.list_external_corpora_communicator.handle_recv,
),
(
ProfileReqOutput,
self.profile_communicator.handle_recv,
@@ -343,6 +370,90 @@ class TokenizerCommunicatorMixin:
]
)
async def add_external_corpus(
self: TokenizerManager, obj: AddExternalCorpusReqInput
) -> AddExternalCorpusReqOutput:
self.auto_create_handle_loop()
truncated = False
try:
if not obj.corpus_id:
import uuid
obj.corpus_id = uuid.uuid4().hex
if obj.file_path is not None:
from sglang.srt.speculative.cpp_ngram.external_corpus import (
iter_external_corpus_chunks,
)
max_tokens = (
self.server_args.speculative_ngram_external_corpus_max_tokens
)
obj.token_chunks = list(
iter_external_corpus_chunks(
obj.file_path, self.tokenizer, max_tokens
)
)
elif obj.documents is not None:
from sglang.srt.speculative.cpp_ngram.external_corpus import (
SEPARATOR_TOKEN,
)
max_tokens = (
self.server_args.speculative_ngram_external_corpus_max_tokens
)
token_chunks = []
total_tokens = 0
has_prev = False
for doc in obj.documents:
if not doc:
continue
token_ids = list(
self.tokenizer.encode(doc, add_special_tokens=False)
)
if not token_ids:
continue
if has_prev:
token_ids = [SEPARATOR_TOKEN] + token_ids
if total_tokens + len(token_ids) > max_tokens:
truncated = True
break
token_chunks.append(token_ids)
total_tokens += len(token_ids)
has_prev = True
obj.token_chunks = token_chunks
else:
return AddExternalCorpusReqOutput(
success=False,
message="Either file_path or documents must be provided.",
)
obj.file_path = None
obj.documents = None
results = await self.add_external_corpus_communicator(obj)
result = results[0]
if truncated and result.success:
result.message += f" (truncated: exceeded {max_tokens} token limit)"
return result
except Exception as e:
return AddExternalCorpusReqOutput(success=False, message=str(e))
async def remove_external_corpus(
self: TokenizerManager, corpus_id: str
) -> RemoveExternalCorpusReqOutput:
self.auto_create_handle_loop()
results = await self.remove_external_corpus_communicator(
RemoveExternalCorpusReqInput(corpus_id=corpus_id)
)
return results[0]
async def list_external_corpora(
self: TokenizerManager,
) -> ListExternalCorporaReqOutput:
self.auto_create_handle_loop()
results = await self.list_external_corpora_communicator(
ListExternalCorporaReqInput()
)
return results[0]
async def flush_cache(
self: TokenizerManager, timeout_s: Optional[float] = None
) -> FlushCacheReqOutput:
python/sglang/srt/server_args.py
+1 -6
View File
@@ -3132,11 +3132,6 @@ class ServerArgs:
"speculative_ngram_external_sam_budget must be less than or equal to "
f"speculative_num_draft_tokens - 1 ({self.speculative_num_draft_tokens - 1})."
)
elif self.speculative_ngram_external_sam_budget != 0:
raise ValueError(
"--speculative-ngram-external-sam-budget requires "
"--speculative-ngram-external-corpus-path."
)
logger.warning(
"The overlap scheduler and mixed chunked prefill are disabled because of "
"using ngram speculative decoding."
@@ -4921,7 +4916,7 @@ class ServerArgs:
"--speculative-ngram-external-corpus-path",
type=str,
default=ServerArgs.speculative_ngram_external_corpus_path,
help="Optional path to an external corpus used to build a read-only SAM for ngram speculative decoding.",
help="Path to an external JSONL corpus to pre-load into SAM at startup. Additional corpora can be added at runtime via POST /add_external_corpus.",
)
parser.add_argument(
"--speculative-ngram-external-sam-budget",
python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py
+12 -48
View File
@@ -1,44 +1,16 @@
# -*- coding: utf-8 -*-
import logging
from collections.abc import Iterable, Iterator, Sequence
from typing import Dict, List, Optional, Tuple
from collections.abc import Iterable, Sequence
from typing import Dict, List, Tuple
import numpy as np
from sglang.jit_kernel.ngram_corpus import get_ngram_corpus_cls
from sglang.srt.speculative.cpp_ngram.external_corpus import SEPARATOR_TOKEN
logger = logging.getLogger(__name__)
# Convenience path for pre-tokenized in-memory documents. The main serving
# startup path loads file-backed corpora through `iter_external_corpus_chunks()`.
def _documents_to_chunks(
documents: Iterable[Sequence[int]],
max_tokens: Optional[int] = None,
) -> Iterator[list[int]]:
total_tokens = 0
has_previous = False
for doc in documents:
if not doc:
continue
doc_tokens = list(doc)
separator_cost = 1 if has_previous else 0
next_total_tokens = total_tokens + separator_cost + len(doc_tokens)
if max_tokens is not None and next_total_tokens > max_tokens:
raise ValueError(
"External ngram corpus exceeds the configured token limit "
f"({max_tokens}) after loading {total_tokens} tokens."
)
total_tokens = next_total_tokens
if has_previous:
yield [SEPARATOR_TOKEN] + doc_tokens
else:
yield doc_tokens
has_previous = True
class NgramCorpus:
def __init__(
self,
@@ -50,7 +22,6 @@ class NgramCorpus:
capacity=1000000,
external_sam_budget=0,
external_corpus_max_tokens=10000000,
external_corpus_documents: Optional[Iterable[Sequence[int]]] = None,
) -> None:
cls = get_ngram_corpus_cls()
self._obj = cls(
@@ -67,13 +38,6 @@ class NgramCorpus:
self.draft_token_num = draft_token_num
self._req_id_to_state_id: Dict[str, int] = {}
self._next_state_id: int = 0
self.external_corpus_token_count = 0
if external_corpus_documents is not None:
self.load_external_corpus(
_documents_to_chunks(
external_corpus_documents, external_corpus_max_tokens
)
)
def _get_state_id(self, req_id: str) -> int:
sid = self._req_id_to_state_id.get(req_id)
@@ -89,18 +53,18 @@ class NgramCorpus:
def synchronize(self):
self._obj.synchronize() # type: ignore
def load_external_corpus(self, chunks: Iterable[Sequence[int]]) -> int:
"""Load pre-chunked external corpus tokens.
Callers passing raw chunk iterables must enforce any token budget before
calling this method. Python-side helpers such as
`iter_external_corpus_chunks()` and `external_corpus_documents=` handle
`external_corpus_max_tokens` validation before handing chunks to C++.
"""
_, loaded_token_count = self._obj.load_external_corpus(chunks)
self.external_corpus_token_count = loaded_token_count
def load_external_corpus_named(
self, corpus_id: str, chunks: Iterable[Sequence[int]]
) -> int:
_, loaded_token_count = self._obj.load_external_corpus_named(corpus_id, chunks)
return loaded_token_count
def remove_external_corpus(self, corpus_id: str) -> None:
self._obj.remove_corpus(corpus_id)
def list_external_corpora(self) -> List[str]:
return self._obj.list_corpora()
def reset(self):
self._obj.reset() # type: ignore
self._req_id_to_state_id.clear()
python/sglang/srt/speculative/external_corpus_manager.py
+103
View File
@@ -0,0 +1,103 @@
"""Manages external SAM corpora for ngram speculative decoding.
Handles add/remove/list operations and async background loading.
Used by the Scheduler — not a mixin, a standalone manager object.
"""
import logging
import threading
from typing import Callable, Optional, Tuple
from sglang.srt.managers.io_struct import (
AddExternalCorpusReqInput,
AddExternalCorpusReqOutput,
ListExternalCorporaReqInput,
ListExternalCorporaReqOutput,
RemoveExternalCorpusReqInput,
RemoveExternalCorpusReqOutput,
)
logger = logging.getLogger(__name__)
class ExternalCorpusManager:
"""Manages external SAM corpus lifecycle for a single scheduler.
Args:
draft_worker: the NGRAMWorker instance (must have add_external_corpus,
remove_external_corpus, list_external_corpora methods).
send_response: callable(output, recv_req) to send deferred responses
back to the tokenizer manager.
"""
def __init__(self, draft_worker, send_response: Callable):
self._worker = draft_worker
self._send_response = send_response
self._pending_load: Optional[
Tuple[AddExternalCorpusReqInput, threading.Thread]
] = None
self._load_result: Optional[AddExternalCorpusReqOutput] = None
def check_pending_load(self):
"""Poll from the scheduler event loop. Sends response when done."""
if self._pending_load is None:
return
recv_req, thread = self._pending_load
if thread.is_alive():
return
self._pending_load = None
thread.join() # formal happens-before for _load_result visibility
result = self._load_result
self._load_result = None
self._send_response(result, recv_req)
def add(
self, recv_req: AddExternalCorpusReqInput
) -> Optional[AddExternalCorpusReqOutput]:
if self._pending_load is not None:
return AddExternalCorpusReqOutput(
success=False,
message="Another corpus load is already in progress.",
)
def _build():
try:
loaded = self._worker.add_external_corpus(
recv_req.corpus_id, recv_req.token_chunks
)
self._load_result = AddExternalCorpusReqOutput(
success=True,
corpus_id=recv_req.corpus_id,
message=f"Loaded corpus '{recv_req.corpus_id}' with {loaded} tokens.",
loaded_token_count=loaded,
)
except Exception as e:
self._load_result = AddExternalCorpusReqOutput(
success=False, message=str(e)
)
thread = threading.Thread(target=_build, daemon=True)
self._pending_load = (recv_req, thread)
thread.start()
return None # response sent later by check_pending_load
def remove(
self, recv_req: RemoveExternalCorpusReqInput
) -> RemoveExternalCorpusReqOutput:
try:
self._worker.remove_external_corpus(recv_req.corpus_id)
return RemoveExternalCorpusReqOutput(
success=True,
message=f"Removed corpus '{recv_req.corpus_id}'.",
)
except Exception as e:
return RemoveExternalCorpusReqOutput(success=False, message=str(e))
def list(
self, recv_req: ListExternalCorporaReqInput
) -> ListExternalCorporaReqOutput:
try:
ids = self._worker.list_external_corpora()
return ListExternalCorporaReqOutput(success=True, corpus_ids=ids)
except Exception as e:
return ListExternalCorporaReqOutput(success=False, message=str(e))
python/sglang/srt/speculative/ngram_worker.py
+20 -7
View File
@@ -11,9 +11,6 @@ from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.server_args import ServerArgs
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.srt.speculative.ngram_info import NgramVerifyInput
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -61,21 +58,37 @@ class NGRAMWorker:
external_corpus_max_tokens=server_args.speculative_ngram_external_corpus_max_tokens,
)
if server_args.speculative_ngram_external_corpus_path is not None:
loaded_token_count = self.ngram_corpus.load_external_corpus(
from sglang.srt.speculative.cpp_ngram.external_corpus import (
iter_external_corpus_chunks,
)
corpus_path = server_args.speculative_ngram_external_corpus_path
chunks = list(
iter_external_corpus_chunks(
server_args.speculative_ngram_external_corpus_path,
corpus_path,
target_worker.tokenizer,
server_args.speculative_ngram_external_corpus_max_tokens,
)
)
loaded = self.add_external_corpus(corpus_path, chunks)
logger.info(
"Loaded external ngram corpus (%d tokens) for SAM speculative decoding.",
loaded_token_count,
"Loaded external ngram corpus '%s' (%d tokens).",
corpus_path,
loaded,
)
def clear_cache_pool(self):
self.ngram_corpus.reset()
def add_external_corpus(self, corpus_id: str, token_chunks: list[list[int]]) -> int:
return self.ngram_corpus.load_external_corpus_named(corpus_id, token_chunks)
def remove_external_corpus(self, corpus_id: str) -> None:
self.ngram_corpus.remove_external_corpus(corpus_id)
def list_external_corpora(self) -> list[str]:
return self.ngram_corpus.list_external_corpora()
def _efficient_concat_last_n(self, seq1: List[int], seq2: List[int], n: int):
seq2_len = len(seq2)
if seq2_len >= n:
test/registered/unit/server_args/test_server_args.py
+8 -15
View File
@@ -427,15 +427,6 @@ class TestHiCacheArgs(unittest.TestCase):
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(
[
@@ -458,12 +449,14 @@ class TestNgramExternalSamArgs(CustomTestCase):
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 _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_external_sam_budget_must_fit_draft_budget(self):
with self.assertRaises(ValueError) as context:
test/registered/unit/spec/test_ngram_corpus.py
+171 -15
View File
@@ -17,6 +17,7 @@ register_cpu_ci(est_time=10, suite="stage-a-test-cpu")
def _make_corpus(match_type="BFS", **kwargs):
external_corpus_documents = kwargs.pop("external_corpus_documents", None)
defaults = dict(
max_trie_depth=12,
min_bfs_breadth=1,
@@ -25,11 +26,23 @@ def _make_corpus(match_type="BFS", **kwargs):
capacity=100000,
external_sam_budget=0,
external_corpus_max_tokens=10000000,
external_corpus_documents=None,
)
defaults.update(kwargs)
defaults["match_type"] = match_type
return NgramCorpus(**defaults)
corpus = NgramCorpus(**defaults)
if external_corpus_documents is not None:
from sglang.srt.speculative.cpp_ngram.external_corpus import SEPARATOR_TOKEN
chunks = []
has_prev = False
for doc in external_corpus_documents:
if has_prev:
chunks.append([SEPARATOR_TOKEN] + list(doc))
else:
chunks.append(list(doc))
has_prev = True
corpus.load_external_corpus_named("test_corpus", chunks)
return corpus
def _batch_get(
@@ -707,8 +720,9 @@ class TestNgramCorpusExternalSam(CustomTestCase):
path = f.name
self.addCleanup(os.remove, path)
loaded_token_count = corpus.load_external_corpus(
iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=8)
loaded_token_count = corpus.load_external_corpus_named(
path,
iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=8),
)
# 5 doc tokens + 1 separator + 2 doc tokens = 8
self.assertEqual(loaded_token_count, 8)
@@ -733,17 +747,9 @@ class TestNgramCorpusExternalSam(CustomTestCase):
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]],
corpus.load_external_corpus_named(
path,
iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=4),
)
def test_external_sam_only_chain(self):
@@ -906,5 +912,155 @@ class TestNgramCorpusMatchBenchmark(CustomTestCase):
)
class TestNgramCorpusMultiSam(CustomTestCase):
"""Verify multi-SAM add/remove/list and budget splitting."""
def test_add_and_list(self):
corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
corpus.load_external_corpus_named("a", [[1, 2, 3, 4, 5]])
corpus.load_external_corpus_named("b", [[10, 20, 30, 40, 50]])
ids = corpus.list_external_corpora()
self.assertEqual(sorted(ids), ["a", "b"])
def test_remove(self):
corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
corpus.load_external_corpus_named("a", [[1, 2, 3, 4, 5]])
corpus.load_external_corpus_named("b", [[10, 20, 30, 40, 50]])
corpus.remove_external_corpus("a")
self.assertEqual(corpus.list_external_corpora(), ["b"])
def test_remove_nonexistent_is_noop(self):
corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
corpus.remove_external_corpus("nonexistent")
self.assertEqual(corpus.list_external_corpora(), [])
def test_multi_sam_candidates(self):
corpus = _make_corpus("BFS", draft_token_num=6, external_sam_budget=4)
corpus.load_external_corpus_named("a", [[1, 2, 3, 10, 11]])
corpus.load_external_corpus_named("b", [[1, 2, 3, 20, 21]])
ids, masks = _batch_get(corpus, [[1, 2, 3]])
leaf_paths = corpus.leaf_paths_from_mask(
ids.tolist(), masks.reshape(6, 6).tolist()
)
# Both SAMs should contribute candidates
self.assertIn([3, 10, 11], leaf_paths)
self.assertIn([3, 20, 21], leaf_paths)
def test_remove_reduces_candidates(self):
corpus = _make_corpus("BFS", draft_token_num=6, external_sam_budget=4)
corpus.load_external_corpus_named("a", [[1, 2, 3, 10, 11]])
corpus.load_external_corpus_named("b", [[1, 2, 3, 20, 21]])
corpus.remove_external_corpus("b")
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.assertNotIn([3, 20, 21], leaf_paths)
def test_make_corpus_with_documents(self):
"""_make_corpus helper loads documents as a named corpus."""
corpus = _make_corpus(
"BFS",
draft_token_num=4,
external_sam_budget=3,
external_corpus_documents=[[1, 2, 3, 4, 5]],
)
ids = corpus.list_external_corpora()
self.assertIn("test_corpus", ids)
class TestMultiSamHttpMock(CustomTestCase):
"""Test HTTP endpoints for multi-SAM management with a mocked backend."""
@classmethod
def setUpClass(cls):
from unittest.mock import AsyncMock, MagicMock
try:
from starlette.testclient import TestClient
from sglang.srt.entrypoints.http_server import app, set_global_state
except (ImportError, OSError):
raise unittest.SkipTest(
"http_server import requires CUDA libraries not available on CPU"
)
from sglang.srt.managers.io_struct import (
AddExternalCorpusReqOutput,
ListExternalCorporaReqOutput,
RemoveExternalCorpusReqOutput,
)
mock_state = MagicMock()
tm = mock_state.tokenizer_manager
# Wire up async methods that the HTTP handlers call
tm.add_external_corpus = AsyncMock(
return_value=AddExternalCorpusReqOutput(
success=True,
corpus_id="test-id",
message="Loaded corpus 'test-id' with 100 tokens.",
loaded_token_count=100,
)
)
tm.remove_external_corpus = AsyncMock(
return_value=RemoveExternalCorpusReqOutput(
success=True, message="Removed corpus 'test-id'."
)
)
tm.list_external_corpora = AsyncMock(
return_value=ListExternalCorporaReqOutput(
success=True, corpus_ids=["a", "b"]
)
)
set_global_state(mock_state)
cls.client = TestClient(app)
cls.mock_tm = tm
def test_add_corpus(self):
resp = self.client.post(
"/add_external_corpus",
json={"corpus_id": "my-corpus", "documents": ["hello world"]},
)
self.assertEqual(resp.status_code, 200)
data = resp.json()
self.assertTrue(data["success"])
self.assertEqual(data["corpus_id"], "test-id")
self.assertEqual(data["loaded_token_count"], 100)
def test_add_corpus_auto_id(self):
resp = self.client.post(
"/add_external_corpus",
json={"documents": ["hello world"]},
)
self.assertEqual(resp.status_code, 200)
self.assertTrue(resp.json()["success"])
def test_remove_corpus(self):
resp = self.client.post(
"/remove_external_corpus",
json={"corpus_id": "test-id"},
)
self.assertEqual(resp.status_code, 200)
self.assertTrue(resp.json()["success"])
def test_remove_corpus_missing_id(self):
resp = self.client.post(
"/remove_external_corpus",
json={},
)
self.assertEqual(resp.status_code, 400)
def test_list_corpora(self):
resp = self.client.get("/list_external_corpora")
self.assertEqual(resp.status_code, 200)
data = resp.json()
self.assertTrue(data["success"])
self.assertEqual(sorted(data["corpus_ids"]), ["a", "b"])
if __name__ == "__main__":
unittest.main(verbosity=3)