diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp index 192b48f6a..1f48aaddd 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp @@ -63,27 +63,57 @@ void Ngram::asyncInsert(std::vector>&& 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 lock(mutex_); - sam_ = std::make_unique(); + if (staging_sam_) { + throw std::runtime_error("startExternalCorpusLoad called while another load is in progress"); + } + staging_sam_ = std::make_unique(); } void Ngram::appendExternalCorpusTokens(const std::vector& tokens) { - std::unique_lock 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 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 lock(mutex_); + sams_[corpus_id] = std::move(staging_sam_); +} + +void Ngram::removeExternalCorpus(const std::string& corpus_id) { + std::unique_lock lock(mutex_); + sams_.erase(corpus_id); } void Ngram::clearExternalCorpus() { std::unique_lock lock(mutex_); - sam_.reset(); + sams_.clear(); + staging_sam_.reset(); +} + +std::vector Ngram::listExternalCorpora() const { + std::unique_lock lock(mutex_); + std::vector 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(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(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; } diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h index fa94d15af..ee77d7d7c 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h @@ -19,12 +19,16 @@ namespace ngram { class Ngram { std::unique_ptr trie_; - std::unique_ptr sam_; + std::unordered_map> 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 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& tokens); - void finishExternalCorpusLoad(); + void finishExternalCorpusLoad(const std::string& corpus_id); + + void removeExternalCorpus(const std::string& corpus_id); void clearExternalCorpus(); + std::vector listExternalCorpora() const; + Result batchMatch(const std::vector>& tokens); Result batchMatch( diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp index 02f4c3916..4f30d5c06 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp @@ -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); } diff --git a/python/sglang/jit_kernel/ngram_corpus.py b/python/sglang/jit_kernel/ngram_corpus.py index 2a6ba13e8..760679cf1 100644 --- a/python/sglang/jit_kernel/ngram_corpus.py +++ b/python/sglang/jit_kernel/ngram_corpus.py @@ -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 diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 6978e0c06..61f7513f5 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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(): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index bd9796534..a36f0ffb1 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 060fdbe4b..bdf1010dd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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]: diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 544c60940..8ba0fe077 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 57dbc1998..f7e0d8ee9 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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", diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py index b7d5f2a9b..5f5433a3d 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py +++ b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py @@ -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() diff --git a/python/sglang/srt/speculative/external_corpus_manager.py b/python/sglang/srt/speculative/external_corpus_manager.py new file mode 100644 index 000000000..5c2d1ec33 --- /dev/null +++ b/python/sglang/srt/speculative/external_corpus_manager.py @@ -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)) diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index afc761a2d..8e133721f 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -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: diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index d18c5ab99..a83a2c693 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -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: diff --git a/test/registered/unit/spec/test_ngram_corpus.py b/test/registered/unit/spec/test_ngram_corpus.py index 2e652f260..4478414aa 100644 --- a/test/registered/unit/spec/test_ngram_corpus.py +++ b/test/registered/unit/spec/test_ngram_corpus.py @@ -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)