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

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

@@ -63,6 +63,29 @@ void Ngram::asyncInsert(std::vector<std::vector<int32_t>>&& tokens) {
   }
 }
 
+void Ngram::startExternalCorpusLoad() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  sam_ = std::make_unique<SuffixAutomaton>();
+}
+
+void Ngram::appendExternalCorpusTokens(const std::vector<int32_t>& tokens) {
+  std::unique_lock<std::mutex> lock(mutex_);
+  sam_->appendTokens(tokens);
+}
+
+void Ngram::finishExternalCorpusLoad() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  sam_->finalize();
+  if (sam_->empty()) {
+    sam_.reset();
+  }
+}
+
+void Ngram::clearExternalCorpus() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  sam_.reset();
+}
+
 void Ngram::insertWorker() {
   for (;;) {
     std::vector<int32_t> data;
@@ -116,12 +139,17 @@ Result Ngram::batchMatch(
 
   std::unique_lock<std::mutex> lock(mutex_);
 
-  using BuildFn = Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&, MatchState&, size_t) const;
-  BuildFn build_fn;
+  using TrieResultBuildFn =
+      Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&, MatchState&, size_t) const;
+  using SamResultBuildFn = Result (SuffixAutomaton::*)(const int32_t*, size_t, int32_t, size_t, const Param&) const;
+  TrieResultBuildFn trie_result_build_fn;
+  SamResultBuildFn sam_result_build_fn;
   if (param_.match_type == "BFS") {
-    build_fn = &Trie::buildRecency;
+    trie_result_build_fn = &Trie::buildRecency;
+    sam_result_build_fn = &SuffixAutomaton::buildRecency;
   } else if (param_.match_type == "PROB") {
-    build_fn = &Trie::buildFrequency;
+    trie_result_build_fn = &Trie::buildFrequency;
+    sam_result_build_fn = &SuffixAutomaton::buildFrequency;
   } else {
     throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'.");
   }
@@ -134,9 +162,23 @@ Result Ngram::batchMatch(
     }
 
     auto& state = match_state_[state_ids[i]];
-    auto draft_token_num = param_.get_draft_token_num(tokens.size());
-    auto res = (trie_.get()->*build_fn)(
-        suffix.data(), suffix.size(), suffix.back(), draft_token_num, param_, state, total_lens[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) {
+      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());
+      merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());
+      continue;
+    }
+
+    auto trie_res = (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());
   }

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

@@ -3,6 +3,7 @@
 #include "param.h"
 #include "queue.h"
 #include "result.h"
+#include "suffix_automaton.h"
 #include "trie.h"
 #include <condition_variable>
 #include <cstddef>
@@ -18,6 +19,7 @@ namespace ngram {
 
 class Ngram {
   std::unique_ptr<Trie> trie_;
+  std::unique_ptr<SuffixAutomaton> sam_;
   Param param_;
 
   // NOTE: protects trie_ and pending_count_. Ensures batchMatch never reads
@@ -40,6 +42,14 @@ class Ngram {
 
   void asyncInsert(std::vector<std::vector<int32_t>>&& tokens);
 
+  void startExternalCorpusLoad();
+
+  void appendExternalCorpusTokens(const std::vector<int32_t>& tokens);
+
+  void finishExternalCorpusLoad();
+
+  void clearExternalCorpus();
+
   Result batchMatch(const std::vector<std::vector<int32_t>>& tokens);
 
   Result batchMatch(

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

@@ -23,7 +23,9 @@ struct NgramCorpusObj : public tvm::ffi::Object {
       int64_t min_bfs_breadth,
       int64_t max_bfs_breadth,
       int64_t draft_token_num,
-      int64_t match_type) {
+      int64_t match_type,
+      int64_t external_sam_budget,
+      int64_t external_corpus_max_tokens) {
     ngram::Param param;
     param.enable = true;
     param.enable_router_mode = false;
@@ -32,6 +34,8 @@ struct NgramCorpusObj : public tvm::ffi::Object {
     param.max_bfs_breadth = static_cast<size_t>(max_bfs_breadth);
     param.draft_token_num = static_cast<size_t>(draft_token_num);
     param.match_type = (match_type == 0) ? "BFS" : "PROB";
+    param.external_sam_budget = static_cast<size_t>(external_sam_budget);
+    param.external_corpus_max_tokens = static_cast<size_t>(external_corpus_max_tokens);
     ngram_ = std::make_unique<ngram::Ngram>(static_cast<size_t>(capacity), param);
   }
 
@@ -97,6 +101,25 @@ struct NgramCorpusObj : public tvm::ffi::Object {
     ngram_->eraseMatchState(state_ids);
   }
 
+  void start_external_corpus_load() {
+    ngram_->startExternalCorpusLoad();
+  }
+
+  void append_external_corpus_tokens(const tvm::ffi::TensorView tokens_tv) {
+    auto* data = static_cast<const int32_t*>(tokens_tv.data_ptr());
+    int64_t n = tokens_tv.size(0);
+    std::vector<int32_t> tokens(data, data + n);
+    ngram_->appendExternalCorpusTokens(tokens);
+  }
+
+  void finish_external_corpus_load() {
+    ngram_->finishExternalCorpusLoad();
+  }
+
+  void clear_external_corpus() {
+    ngram_->clearExternalCorpus();
+  }
+
   void synchronize() {
     ngram_->synchronize();
   }
@@ -130,11 +153,15 @@ struct NgramCorpusObj : public tvm::ffi::Object {
 void register_ngram_corpus() {
   namespace refl = tvm::ffi::reflection;
   refl::ObjectDef<NgramCorpusObj>()
-      .def(refl::init<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>(), "__init__")
+      .def(refl::init<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>(), "__init__")
       .def("async_insert", &NgramCorpusObj::async_insert)
       .def("batch_match", &NgramCorpusObj::batch_match)
       .def("batch_match_stateful", &NgramCorpusObj::batch_match_stateful)
       .def("erase_match_state", &NgramCorpusObj::erase_match_state)
+      .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("clear_external_corpus", &NgramCorpusObj::clear_external_corpus)
       .def("synchronize", &NgramCorpusObj::synchronize)
       .def("reset", &NgramCorpusObj::reset);
 }

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

@@ -19,6 +19,8 @@ struct Param {
   size_t max_bfs_breadth;
   size_t max_trie_depth;
   size_t draft_token_num;
+  size_t external_sam_budget = 0;
+  size_t external_corpus_max_tokens = 10000000;
   std::string match_type;
 
   std::vector<size_t> batch_draft_token_num;
@@ -92,7 +94,8 @@ struct Param {
     ss << "enable = " << enable << ", enable_router_mode = " << enable_router_mode
        << ", min_bfs_breadth = " << min_bfs_breadth << ", max_bfs_breadth = " << max_bfs_breadth
        << ", max_trie_depth = " << max_trie_depth << ", draft_token_num = " << draft_token_num
-       << ", match_type = " << match_type;
+       << ", external_sam_budget = " << external_sam_budget
+       << ", external_corpus_max_tokens = " << external_corpus_max_tokens << ", match_type = " << match_type;
     ss << ", batch_draft_token_num(" << batch_draft_token_num.size() << ") = ";
     for (int i = 0; i < batch_draft_token_num.size(); ++i) {
       ss << i << "|" << batch_draft_token_num[i] << ",";

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

@@ -1,5 +1,6 @@
 #include "result.h"
 
+#include <algorithm>
 #include <cstring>
 #include <queue>
 #include <tuple>
@@ -47,6 +48,81 @@ Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree,
   return info;
 }
 
+std::vector<std::vector<int32_t>> extractLeafPaths_(const Result& result) {
+  const auto n = static_cast<int>(result.token.size());
+  if (n <= 1) {
+    return {};
+  }
+
+  std::vector<int> parent(n, -1);
+  std::vector<bool> has_child(n, false);
+  for (int i = 1; i < n; ++i) {
+    for (int j = i - 1; j >= 0; --j) {
+      if (result.mask[i * n + j]) {
+        parent[i] = j;
+        has_child[j] = true;
+        break;
+      }
+    }
+  }
+
+  std::vector<std::vector<int32_t>> paths;
+  for (int leaf = 1; leaf < n; ++leaf) {
+    if (has_child[leaf]) {
+      continue;
+    }
+    std::vector<int32_t> path;
+    for (int cursor = leaf; cursor > 0; cursor = parent[cursor]) {
+      path.emplace_back(result.token[cursor]);
+    }
+    std::reverse(path.begin(), path.end());
+    if (path.size() == 1 && path.front() == 0) {
+      continue;
+    }
+    paths.emplace_back(std::move(path));
+  }
+  return paths;
+}
+
+Result buildResultFromLeafPaths_(int last_token, int draft_token_num, const std::vector<std::vector<int32_t>>& paths) {
+  std::vector<Node> tree(draft_token_num);
+  const int root = 0;
+  int cursor = 1;
+  for (const auto& path : paths) {
+    int parent = root;
+    for (const auto token : path) {
+      auto iter = tree[parent].next.find(token);
+      if (iter == tree[parent].next.end()) {
+        if (cursor >= draft_token_num) {
+          parent = -1;
+          break;
+        }
+        iter = tree[parent].next.insert({token, cursor++}).first;
+      }
+      parent = iter->second;
+    }
+    if (cursor >= draft_token_num) {
+      break;
+    }
+  }
+  return fillResult(last_token, draft_token_num, tree, root);
+}
+
+Result combineRootResults_(int last_token, int draft_token_num, const Result& primary, const Result& secondary) {
+  auto primary_paths = extractLeafPaths_(primary);
+  auto secondary_paths = extractLeafPaths_(secondary);
+  std::vector<std::vector<int32_t>> merged_paths = std::move(primary_paths);
+  merged_paths.reserve(merged_paths.size() + secondary_paths.size());
+  for (const auto& path : secondary_paths) {
+    if (path.empty()) {
+      continue;
+    }
+    merged_paths.emplace_back(path);
+  }
+
+  return buildResultFromLeafPaths_(last_token, draft_token_num, merged_paths);
+}
+
 void Result::truncate(size_t n) {
   if (n < token.size()) {
     int full_n = token.size();

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

@@ -18,5 +18,8 @@ struct Node {
 };
 
 Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root);
+std::vector<std::vector<int32_t>> extractLeafPaths_(const Result& result);
+Result buildResultFromLeafPaths_(int last_token, int draft_token_num, const std::vector<std::vector<int32_t>>& paths);
+Result combineRootResults_(int last_token, int draft_token_num, const Result& primary, const Result& secondary);
 
 }  // namespace ngram

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

@@ -0,0 +1,283 @@
+#include "suffix_automaton.h"
+
+#include <algorithm>
+#include <numeric>
+#include <queue>
+#include <stdexcept>
+#include <string>
+#include <tuple>
+
+namespace ngram {
+
+SuffixAutomaton::SuffixAutomaton() {
+  reset_();
+}
+
+void SuffixAutomaton::reset_() {
+  states_.clear();
+  states_.emplace_back();
+  last_ = 0;
+  pos_ = 0;
+  saw_token_ = false;
+  finalized_ = false;
+  loaded_ = false;
+}
+
+void SuffixAutomaton::appendTokens(const std::vector<int32_t>& tokens) {
+  if (finalized_) {
+    throw std::runtime_error("Cannot append tokens after finalizing the SAM.");
+  }
+  if (tokens.empty()) {
+    return;
+  }
+
+  for (const auto token : tokens) {
+    extend_(token, pos_++);
+    saw_token_ = true;
+  }
+}
+
+void SuffixAutomaton::finalize() {
+  if (finalized_) {
+    return;
+  }
+  finalized_ = true;
+  if (!saw_token_) {
+    return;
+  }
+
+  propagateOccurrencesAndRecency_();
+  loaded_ = true;
+}
+
+void SuffixAutomaton::extend_(int32_t token, int64_t pos) {
+  const int cur = static_cast<int>(states_.size());
+  states_.emplace_back();
+  states_[cur].max_len = states_[last_].max_len + 1;
+  states_[cur].occ_count = 1;
+  states_[cur].max_end_pos = pos;
+
+  int p = last_;
+  while (p != -1 && !states_[p].next.contains(token)) {
+    states_[p].next[token] = cur;
+    p = states_[p].link;
+  }
+
+  if (p == -1) {
+    states_[cur].link = 0;
+    last_ = cur;
+    return;
+  }
+
+  const int q = states_[p].next[token];
+  if (states_[p].max_len + 1 == states_[q].max_len) {
+    states_[cur].link = q;
+    last_ = cur;
+    return;
+  }
+
+  const int clone = static_cast<int>(states_.size());
+  states_.push_back(states_[q]);
+  states_[clone].max_len = states_[p].max_len + 1;
+  states_[clone].occ_count = 0;
+  states_[clone].children_by_freq.clear();
+  states_[clone].children_by_recency.clear();
+
+  while (p != -1 && states_[p].next[token] == q) {
+    states_[p].next[token] = clone;
+    p = states_[p].link;
+  }
+
+  states_[q].link = clone;
+  states_[cur].link = clone;
+  last_ = cur;
+}
+
+void SuffixAutomaton::propagateOccurrencesAndRecency_() {
+  std::vector<int> order(states_.size());
+  std::iota(order.begin(), order.end(), 0);
+  std::sort(
+      order.begin(), order.end(), [this](int lhs, int rhs) { return states_[lhs].max_len < states_[rhs].max_len; });
+
+  for (auto it = order.rbegin(); it != order.rend(); ++it) {
+    const int state = *it;
+    const int link = states_[state].link;
+    if (link < 0) {
+      continue;
+    }
+    states_[link].occ_count += states_[state].occ_count;
+    states_[link].max_end_pos = std::max(states_[link].max_end_pos, states_[state].max_end_pos);
+  }
+
+  for (auto& state : states_) {
+    state.children_by_freq.clear();
+    state.children_by_recency.clear();
+    state.children_by_freq.reserve(state.next.size());
+    state.children_by_recency.reserve(state.next.size());
+    for (const auto& [token, child_state] : state.next) {
+      if (token == kSeparatorToken) {
+        continue;
+      }
+      state.children_by_freq.emplace_back(token, child_state);
+      state.children_by_recency.emplace_back(token, child_state);
+    }
+
+    std::sort(state.children_by_freq.begin(), state.children_by_freq.end(), [this](const auto& lhs, const auto& rhs) {
+      const auto lhs_freq = states_[lhs.second].occ_count;
+      const auto rhs_freq = states_[rhs.second].occ_count;
+      return std::tie(rhs_freq, lhs.first, lhs.second) < std::tie(lhs_freq, rhs.first, rhs.second);
+    });
+    std::sort(
+        state.children_by_recency.begin(), state.children_by_recency.end(), [this](const auto& lhs, const auto& rhs) {
+          const auto lhs_recency = states_[lhs.second].max_end_pos;
+          const auto rhs_recency = states_[rhs.second].max_end_pos;
+          return std::tie(rhs_recency, lhs.first, lhs.second) < std::tie(lhs_recency, rhs.first, rhs.second);
+        });
+  }
+}
+
+std::vector<SamAnchor> SuffixAutomaton::match(const int32_t* context, size_t len, size_t max_depth) const {
+  if (empty() || len == 0) {
+    return {};
+  }
+
+  const auto start = len > max_depth ? len - max_depth : 0;
+  int state = 0;
+  int32_t matched_len = 0;
+  for (size_t i = start; i < len; ++i) {
+    const auto token = context[i];
+    while (state != 0 && !states_[state].next.contains(token)) {
+      state = states_[state].link;
+      matched_len = std::min<int32_t>(matched_len, states_[state].max_len);
+    }
+    if (auto iter = states_[state].next.find(token); iter != states_[state].next.end()) {
+      state = iter->second;
+      ++matched_len;
+    } else if (auto root_iter = states_[0].next.find(token); root_iter != states_[0].next.end()) {
+      state = root_iter->second;
+      matched_len = 1;
+    } else {
+      state = 0;
+      matched_len = 0;
+    }
+  }
+
+  std::vector<SamAnchor> anchors;
+  while (state > 0 && matched_len > 0) {
+    if (!states_[state].children_by_freq.empty()) {
+      anchors.push_back({state, matched_len});
+    }
+    state = states_[state].link;
+    if (state <= 0) {
+      break;
+    }
+    matched_len = std::min<int32_t>(matched_len, states_[state].max_len);
+  }
+  return anchors;
+}
+
+Result SuffixAutomaton::buildRecency(
+    const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const {
+  auto anchors = match(context, len, param.max_trie_depth);
+  const auto max_match_depth = std::max<int32_t>(1, static_cast<int32_t>(param.max_trie_depth - 1));
+  const double bfs_breadth_scale = double(param.max_bfs_breadth - param.min_bfs_breadth) / max_match_depth;
+  std::vector<Node> tree(draft_token_num + 1);
+  int root = 0;
+  int cursor = 1;
+
+  for (const auto& anchor : anchors) {
+    std::queue<std::tuple<int, double, int>> queue;
+    queue.push(
+        {root, (max_match_depth - anchor.matched_len) * bfs_breadth_scale + param.min_bfs_breadth, anchor.state});
+    while (!queue.empty() && cursor <= static_cast<int>(draft_token_num)) {
+      auto [parent, cur_breadth, state] = queue.front();
+      queue.pop();
+
+      const auto& children = states_[state].children_by_recency;
+      const auto breadth = std::max(1, static_cast<int32_t>(cur_breadth));
+      for (int i = 0;
+           i < breadth && i < static_cast<int>(children.size()) && cursor <= static_cast<int>(draft_token_num);
+           ++i) {
+        const auto [token, child_state] = children[i];
+        int pos = -1;
+        if (auto iter = tree[parent].next.find(token); iter != tree[parent].next.end()) {
+          pos = iter->second;
+        } else {
+          pos = tree[parent].next.insert({token, cursor++}).first->second;
+        }
+        queue.emplace(pos, cur_breadth - bfs_breadth_scale, child_state);
+      }
+    }
+  }
+  return fillResult(last_token, draft_token_num + 1, tree, root);
+}
+
+Result SuffixAutomaton::buildFrequency(
+    const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const {
+  auto anchors = match(context, len, param.max_trie_depth);
+  struct CompareByProb {
+    bool operator()(
+        const std::tuple<int, int32_t, int, double>& lhs, const std::tuple<int, int32_t, int, double>& rhs) const {
+      return std::get<3>(lhs) < std::get<3>(rhs);
+    }
+  };
+
+  std::priority_queue<
+      std::tuple<int, int32_t, int, double>,
+      std::vector<std::tuple<int, int32_t, int, double>>,
+      CompareByProb>
+      heap;
+  std::vector<Node> tree(draft_token_num + 1);
+  int root = 0;
+  int cursor = 1;
+  const int top_k = static_cast<int>(param.max_bfs_breadth);
+
+  auto addToHeap = [this, &heap, top_k](int parent, int state, double prob) {
+    if (top_k <= 0) {
+      return;
+    }
+    const auto& children = states_[state].children_by_freq;
+    if (children.empty()) {
+      return;
+    }
+    double sum_freq = 0.0;
+    int count = 0;
+    for (const auto& [_, child_state] : children) {
+      sum_freq += static_cast<double>(states_[child_state].occ_count);
+      if (++count >= top_k) {
+        break;
+      }
+    }
+    if (sum_freq <= 0) {
+      sum_freq = 1.0;
+    }
+    count = 0;
+    for (const auto& [token, child_state] : children) {
+      const auto scaled_prob = static_cast<double>(states_[child_state].occ_count) / sum_freq * prob;
+      heap.emplace(parent, token, child_state, scaled_prob);
+      if (++count >= top_k) {
+        break;
+      }
+    }
+  };
+
+  for (const auto& anchor : anchors) {
+    addToHeap(root, anchor.state, 1.0);
+    while (!heap.empty() && cursor <= static_cast<int>(draft_token_num)) {
+      auto [parent, token, child_state, prob] = heap.top();
+      heap.pop();
+
+      int pos = -1;
+      if (auto iter = tree[parent].next.find(token); iter != tree[parent].next.end()) {
+        pos = iter->second;
+      } else {
+        pos = cursor++;
+        tree[parent].next[token] = pos;
+      }
+      addToHeap(pos, child_state, prob);
+    }
+  }
+  return fillResult(last_token, draft_token_num + 1, tree, root);
+}
+
+}  // namespace ngram

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

@@ -0,0 +1,62 @@
+#pragma once
+
+#include "param.h"
+#include "result.h"
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <unordered_map>
+#include <vector>
+
+namespace ngram {
+
+struct SamAnchor {
+  int state = 0;
+  int32_t matched_len = 0;
+};
+
+struct SamState {
+  int link = -1;
+  int32_t max_len = 0;
+  std::unordered_map<int32_t, int> next;
+  uint64_t occ_count = 0;
+  int64_t max_end_pos = -1;
+  std::vector<std::pair<int32_t, int>> children_by_freq;
+  std::vector<std::pair<int32_t, int>> children_by_recency;
+};
+
+class SuffixAutomaton {
+ public:
+  static constexpr int32_t kSeparatorToken = std::numeric_limits<int32_t>::min();
+
+  SuffixAutomaton();
+
+  void appendTokens(const std::vector<int32_t>& tokens);
+
+  void finalize();
+
+  bool empty() const {
+    return !loaded_;
+  }
+
+  Result buildRecency(
+      const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const;
+
+  Result buildFrequency(
+      const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const;
+
+ private:
+  void reset_();
+  void extend_(int32_t token, int64_t pos);
+  void propagateOccurrencesAndRecency_();
+  std::vector<SamAnchor> match(const int32_t* context, size_t len, size_t max_depth) const;
+
+  std::vector<SamState> states_;
+  int last_ = 0;
+  int64_t pos_ = 0;
+  bool saw_token_ = false;
+  bool finalized_ = false;
+  bool loaded_ = false;
+};
+
+}  // namespace ngram

python/sglang/jit_kernel/ngram_corpus.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Iterable, Sequence
 from typing import List, Tuple
 
 import numpy as np
@@ -29,6 +30,7 @@ def get_ngram_corpus_cls():
         cpp_files=[
             "ngram_corpus/result.cpp",
             "ngram_corpus/trie.cpp",
+            "ngram_corpus/suffix_automaton.cpp",
             "ngram_corpus/ngram.cpp",
             "ngram_corpus/ngram_corpus_ffi.cpp",
         ],
@@ -48,6 +50,8 @@ def get_ngram_corpus_cls():
             max_bfs_breadth: int,
             draft_token_num: int,
             match_type: str,
+            external_sam_budget: int = 0,
+            external_corpus_max_tokens: int = 10000000,
         ) -> None:
             mt = _MATCH_TYPE_MAP.get(match_type)
             if mt is None:
@@ -61,6 +65,8 @@ def get_ngram_corpus_cls():
                 max_bfs_breadth,
                 draft_token_num,
                 mt,
+                external_sam_budget,
+                external_corpus_max_tokens,
             )
             self._draft_token_num = draft_token_num
 
@@ -112,4 +118,22 @@ 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]]
+        ) -> Tuple[int, int]:
+            self.start_external_corpus_load()  # type: ignore
+            chunk_count = 0
+            loaded_token_count = 0
+            try:
+                for chunk in chunks:
+                    tokens_t = torch.tensor(list(chunk), dtype=torch.int32)
+                    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
+            except Exception:
+                self.clear_external_corpus()  # type: ignore
+                raise
+            return chunk_count, loaded_token_count
+
     return NgramCorpusFFI

python/sglang/srt/server_args.py

@@ -512,6 +512,9 @@ class ServerArgs:
     speculative_ngram_match_type: Literal["BFS", "PROB"] = "BFS"
     speculative_ngram_max_trie_depth: int = 18
     speculative_ngram_capacity: int = 10 * 1000 * 1000
+    speculative_ngram_external_corpus_path: Optional[str] = None
+    speculative_ngram_external_sam_budget: int = 0
+    speculative_ngram_external_corpus_max_tokens: int = 10000000
     enable_multi_layer_eagle: bool = False
 
     # Expert parallelism
@@ -3110,6 +3113,30 @@ class ServerArgs:
                     "speculative_num_draft_tokens is set to 12 by default for ngram speculative decoding. "
                     "You can override this by explicitly setting --speculative-num-draft-tokens."
                 )
+            if self.speculative_ngram_external_corpus_path is not None:
+                if self.speculative_ngram_external_sam_budget <= 0:
+                    raise ValueError(
+                        "--speculative-ngram-external-sam-budget must be positive when "
+                        "--speculative-ngram-external-corpus-path is set."
+                    )
+                if self.speculative_ngram_external_corpus_max_tokens <= 0:
+                    raise ValueError(
+                        "--speculative-ngram-external-corpus-max-tokens must be positive when "
+                        "--speculative-ngram-external-corpus-path is set."
+                    )
+                if (
+                    self.speculative_ngram_external_sam_budget
+                    > self.speculative_num_draft_tokens - 1
+                ):
+                    raise ValueError(
+                        "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."
@@ -4888,6 +4915,24 @@ class ServerArgs:
             default=ServerArgs.speculative_ngram_capacity,
             help="The cache capacity for ngram speculative decoding.",
         )
+        parser.add_argument(
+            "--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.",
+        )
+        parser.add_argument(
+            "--speculative-ngram-external-sam-budget",
+            type=int,
+            default=ServerArgs.speculative_ngram_external_sam_budget,
+            help="Number of draft nodes reserved for the external SAM subtree in ngram speculative decoding.",
+        )
+        parser.add_argument(
+            "--speculative-ngram-external-corpus-max-tokens",
+            type=int,
+            default=ServerArgs.speculative_ngram_external_corpus_max_tokens,
+            help="Fail startup if the tokenized external ngram corpus exceeds this many tokens. Tune this based on your CPU memory budget.",
+        )
 
         # Multi-layer Eagle speculative decoding
         parser.add_argument(

python/sglang/srt/speculative/cpp_ngram/external_corpus.py

@@ -0,0 +1,62 @@
+import json
+from collections.abc import Iterator
+from pathlib import Path
+
+# Must match SuffixAutomaton::kSeparatorToken in suffix_automaton.h.
+SEPARATOR_TOKEN = -(2**31)
+
+# Default chunk size for streaming tokenized documents into the SAM.
+DEFAULT_CHUNK_SIZE = 4096
+
+
+def iter_external_corpus_chunks(
+    path: str, tokenizer, max_tokens: int, chunk_size: int = DEFAULT_CHUNK_SIZE
+) -> Iterator[list[int]]:
+    """Chunk documents and yield fixed-size token chunks from a JSONL corpus file."""
+    corpus_path = Path(path)
+    if not corpus_path.is_file():
+        raise ValueError(f"External ngram corpus path does not exist: {path}")
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required to load an external ngram corpus.")
+    if max_tokens <= 0:
+        raise ValueError("External ngram corpus max tokens must be positive.")
+
+    total_tokens = 0
+    has_previous_doc = False
+    with corpus_path.open("r", encoding="utf-8") as f:
+        for line_no, line in enumerate(f, start=1):
+            if not line.strip():
+                continue
+
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"Invalid JSON in external ngram corpus at line {line_no}: {e.msg}"
+                ) from e
+
+            if not isinstance(record, str):
+                raise ValueError(
+                    "Invalid external ngram corpus record at line "
+                    f"{line_no}: expected a JSON string."
+                )
+
+            token_ids = list(tokenizer.encode(record, add_special_tokens=False))
+            if not token_ids:
+                continue
+
+            separator_cost = 1 if has_previous_doc else 0
+            next_total_tokens = total_tokens + separator_cost + len(token_ids)
+            if next_total_tokens > max_tokens:
+                raise ValueError(
+                    "External ngram corpus exceeds the configured token limit "
+                    f"({max_tokens}) at line {line_no} after loading "
+                    f"{total_tokens} tokens."
+                )
+            total_tokens = next_total_tokens
+
+            if has_previous_doc:
+                token_ids = [SEPARATOR_TOKEN] + token_ids
+            for i in range(0, len(token_ids), chunk_size):
+                yield token_ids[i : i + chunk_size]
+            has_previous_doc = True

python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py

@@ -1,15 +1,44 @@
 # -*- coding: utf-8 -*-
 
 import logging
-from typing import Dict, List, Tuple
+from collections.abc import Iterable, Iterator, Sequence
+from typing import Dict, List, Optional, 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,
@@ -19,6 +48,9 @@ class NgramCorpus:
         draft_token_num=8,
         match_type="BFS",
         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(
@@ -28,11 +60,20 @@ class NgramCorpus:
             max_bfs_breadth=max_bfs_breadth,
             draft_token_num=draft_token_num,
             match_type=match_type,
+            external_sam_budget=external_sam_budget,
+            external_corpus_max_tokens=external_corpus_max_tokens,
         )
         self.default_mask = np.ones((1, 1), dtype=np.int64)
         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)
@@ -48,6 +89,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
+        return loaded_token_count
+
     def reset(self):
         self._obj.reset()  # type: ignore
         self._req_id_to_state_id.clear()

python/sglang/srt/speculative/ngram_worker.py

@@ -11,6 +11,9 @@ 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
@@ -54,7 +57,21 @@ class NGRAMWorker:
             capacity=server_args.speculative_ngram_capacity,
             max_trie_depth=server_args.speculative_ngram_max_trie_depth,
             draft_token_num=server_args.speculative_num_draft_tokens,
+            external_sam_budget=server_args.speculative_ngram_external_sam_budget,
+            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(
+                iter_external_corpus_chunks(
+                    server_args.speculative_ngram_external_corpus_path,
+                    target_worker.tokenizer,
+                    server_args.speculative_ngram_external_corpus_max_tokens,
+                )
+            )
+            logger.info(
+                "Loaded external ngram corpus (%d tokens) for SAM speculative decoding.",
+                loaded_token_count,
+            )
 
     def clear_cache_pool(self):
         self.ngram_corpus.reset()