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

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp

@@ -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_);
+  if (staging_sam_) {
-  sam_ = std::make_unique<SuffixAutomaton>();
+    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_);
+  if (!staging_sam_) {
-  sam_->appendTokens(tokens);
+    throw std::runtime_error("appendExternalCorpusTokens called without startExternalCorpusLoad");
+  }
+  staging_sam_->appendTokens(tokens);
 }
 
-void Ngram::finishExternalCorpusLoad() {
+void Ngram::finishExternalCorpusLoad(const std::string& corpus_id) {
-  std::unique_lock<std::mutex> lock(mutex_);
+  if (!staging_sam_) {
-  sam_->finalize();
+    throw std::runtime_error("finishExternalCorpusLoad called without startExternalCorpusLoad");
-  if (sam_->empty()) {
-    sam_.reset();
   }
+  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);
+    for (const auto& [_, sam] : sams_) {
-    merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
+      auto sam_res =
-    merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());
+          (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

@@ -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
+  // NOTE: protects trie_, sams_, and pending_count_. staging_sam_ is NOT
-  // trie_ while insertWorker is writing. After synchronize(), no pending
+  // protected by mutex_ — it is only accessed from the corpus loading thread.
-  // inserts remain so mutex_ contention is effectively zero.
+  // 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

@@ -112,14 +112,28 @@ struct NgramCorpusObj : public tvm::ffi::Object {
     ngram_->appendExternalCorpusTokens(tokens);
   }
 
-  void finish_external_corpus_load() {
+  void finish_external_corpus_load(const std::string& corpus_id) {
-    ngram_->finishExternalCorpusLoad();
+    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

@@ -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(
+        def load_external_corpus_named(
-            self, chunks: Iterable[Sequence[int]]
+            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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -1,44 +1,16 @@
 # -*- coding: utf-8 -*-
 
 import logging
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Iterable, Sequence
-from typing import Dict, List, Optional, Tuple
+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:
+    def load_external_corpus_named(
-        """Load pre-chunked external corpus tokens.
+        self, corpus_id: str, chunks: Iterable[Sequence[int]]
-
+    ) -> int:
-        Callers passing raw chunk iterables must enforce any token budget before
+        _, loaded_token_count = self._obj.load_external_corpus_named(corpus_id, chunks)
-        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
         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

@@ -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

@@ -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 external ngram corpus '%s' (%d tokens).",
-                loaded_token_count,
+                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

@@ -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):
+    def _make_dummy_ngram_args(self, **overrides):
-        with self.assertRaises(ValueError) as context:
+        args = ServerArgs(model_path="dummy")
-            self._make_dummy_ngram_args(
+        args.speculative_algorithm = "NGRAM"
-                speculative_ngram_external_sam_budget=2,
+        args.speculative_num_draft_tokens = 12
-            )._handle_speculative_decoding()
+        args.device = "cuda"
-        self.assertIn("external-sam-budget", str(context.exception))
+        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

@@ -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(
+        loaded_token_count = corpus.load_external_corpus_named(
-            iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=8)
+            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(
+            corpus.load_external_corpus_named(
-                iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=4)
+                path,
-            )
+                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):
@@ -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)