[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