[Spec][Ngram] Misc enhance support for multiple SAMs (#22294)

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

@@ -92,7 +92,11 @@ void Ngram::finishExternalCorpusLoad(const std::string& corpus_id) {
   }
   // Only lock briefly to install the completed SAM.
   std::unique_lock<std::mutex> lock(mutex_);
-  sams_[corpus_id] = std::move(staging_sam_);
+  if (sams_.find(corpus_id) != sams_.end()) {
+    throw std::runtime_error(
+        "External corpus '" + corpus_id + "' already exists. Remove it before adding a new corpus with the same id.");
+  }
+  sams_.emplace(corpus_id, std::move(staging_sam_));
 }
 
 void Ngram::removeExternalCorpus(const std::string& corpus_id) {
@@ -100,6 +104,11 @@ void Ngram::removeExternalCorpus(const std::string& corpus_id) {
   sams_.erase(corpus_id);
 }
 
+void Ngram::resetStagingSam() {
+  // staging_sam_ is only accessed from the loading thread — no lock needed.
+  staging_sam_.reset();
+}
+
 void Ngram::clearExternalCorpus() {
   std::unique_lock<std::mutex> lock(mutex_);
   sams_.clear();

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

@@ -50,10 +50,13 @@ class Ngram {
 
   void appendExternalCorpusTokens(const std::vector<int32_t>& tokens);
 
+  // Publishes the staged corpus. Duplicate corpus_id is rejected.
   void finishExternalCorpusLoad(const std::string& corpus_id);
 
   void removeExternalCorpus(const std::string& corpus_id);
 
+  void resetStagingSam();
+
   void clearExternalCorpus();
 
   std::vector<std::string> listExternalCorpora() const;
@@ -67,6 +70,9 @@ class Ngram {
 
   void eraseMatchState(const std::vector<int64_t>& state_ids);
 
+  // Resets the online trie and match state but preserves external corpora
+  // (sams_). External corpora are user-managed via add/remove APIs and
+  // should not be affected by cache flushes.
   void reset() {
     std::unique_lock<std::mutex> lock(mutex_);
     if (trie_) {

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

@@ -120,6 +120,10 @@ struct NgramCorpusObj : public tvm::ffi::Object {
     ngram_->removeExternalCorpus(corpus_id);
   }
 
+  void cancel_external_corpus_load() {
+    ngram_->resetStagingSam();
+  }
+
   void clear_external_corpus() {
     ngram_->clearExternalCorpus();
   }
@@ -176,6 +180,7 @@ void register_ngram_corpus() {
       .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("cancel_external_corpus_load", &NgramCorpusObj::cancel_external_corpus_load)
       .def("clear_external_corpus", &NgramCorpusObj::clear_external_corpus)
       .def("list_external_corpora", &NgramCorpusObj::list_external_corpora)
       .def("synchronize", &NgramCorpusObj::synchronize)