Migrate ngram corpus from torch cpp_extension to TVM FFI jit_kernel (#21920)

Co-authored-by: DarkSharpness <2040703891@qq.com>

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

@@ -0,0 +1,103 @@
+#include "ngram.h"
+
+#include "trie.h"
+#include <limits>
+#include <stdexcept>
+#include <string>
+
+namespace ngram {
+
+Ngram::Ngram(size_t capacity, const Param& param) : param_(param) {
+  if (!(param_.max_trie_depth > 1)) {
+    throw std::runtime_error(
+        "param_.max_trie_depth must be greater than 1, current value: " + std::to_string(param_.max_trie_depth));
+  }
+  if (!(param_.min_bfs_breadth > 0)) {
+    throw std::runtime_error(
+        "min_bfs_breadth must be greater than 0, current value: " + std::to_string(param_.min_bfs_breadth));
+  }
+  if (!(param_.min_bfs_breadth <= param_.max_bfs_breadth)) {
+    throw std::runtime_error(
+        "min_bfs_breadth must be less than or equal to max_bfs_breadth, "
+        "current min_bfs_breadth: " +
+        std::to_string(param_.min_bfs_breadth) + ", max_bfs_breadth: " + std::to_string(param_.max_bfs_breadth));
+  }
+  if (!(param_.draft_token_num > 0)) {
+    throw std::runtime_error(
+        "draft_token_num must be greater than 0, current value: " + std::to_string(param_.draft_token_num));
+  }
+  for (auto config : param_.batch_draft_token_num) {
+    if (config != std::numeric_limits<decltype(config)>::max()) {
+      if (!(config <= param_.draft_token_num)) {
+        throw std::runtime_error(
+            "batch_draft_token_num config value " + std::to_string(config) +
+            " must be less than or equal to draft_token_num: " + std::to_string(param_.draft_token_num));
+      }
+    }
+  }
+
+  trie_ = std::make_unique<Trie>(capacity, param_);
+
+  insert_worker_ = std::thread(&Ngram::insertWorker, this);
+}
+
+Ngram::~Ngram() {
+  insert_queue_.close();
+  if (insert_worker_.joinable()) {
+    insert_worker_.join();
+  }
+}
+
+void Ngram::synchronize() const {
+  std::unique_lock<std::mutex> lock(mutex_);
+  sync_cv_.wait(lock, [this] { return pending_count_ == 0; });
+}
+
+void Ngram::asyncInsert(std::vector<std::vector<int32_t>>&& tokens) {
+  {
+    std::lock_guard<std::mutex> lock(mutex_);
+    pending_count_ += tokens.size();
+  }
+  for (auto&& token : tokens) {
+    insert_queue_.enqueue(std::move(token));
+  }
+}
+
+void Ngram::insertWorker() {
+  for (;;) {
+    std::vector<int32_t> data;
+    if (!insert_queue_.dequeue(data)) {
+      break;
+    }
+    std::unique_lock<std::mutex> lock(mutex_);
+    trie_->insert(data.data(), data.size());
+    --pending_count_;
+    lock.unlock();
+    sync_cv_.notify_all();
+  }
+}
+
+Result Ngram::batchMatch(const std::vector<std::vector<int32_t>>& tokens) const {
+  std::unique_lock<std::mutex> lock(mutex_);
+
+  using BuildFn = Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&) const;
+  BuildFn build_fn;
+  if (param_.match_type == "BFS") {
+    build_fn = &Trie::buildRecency;
+  } else if (param_.match_type == "PROB") {
+    build_fn = &Trie::buildFrequency;
+  } else {
+    throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'.");
+  }
+
+  Result merged;
+  for (const auto& suffix : tokens) {
+    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_);
+    merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
+    merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());
+  }
+  return merged;
+}
+
+}  // namespace ngram

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

@@ -0,0 +1,57 @@
+#pragma once
+
+#include "param.h"
+#include "queue.h"
+#include "result.h"
+#include "trie.h"
+#include <condition_variable>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+namespace ngram {
+
+class Ngram {
+  std::unique_ptr<Trie> trie_;
+  Param param_;
+
+  // NOTE: protects trie_ and pending_count_. Ensures batchMatch never reads
+  // trie_ while insertWorker is writing. After synchronize(), no pending
+  // inserts remain so mutex_ contention is effectively zero.
+  mutable std::mutex mutex_;
+  mutable std::condition_variable sync_cv_;
+  // NOTE: tracks inserts from enqueue through trie_->insert() completion,
+  // not just queue occupancy. A dequeued item may still be mid-insert.
+  size_t pending_count_ = 0;
+  utils::Queue<std::vector<int32_t>> insert_queue_;
+  std::thread insert_worker_;
+
+ public:
+  Ngram(size_t capacity, const Param& param);
+  ~Ngram();
+
+  void synchronize() const;
+
+  void asyncInsert(std::vector<std::vector<int32_t>>&& tokens);
+
+  Result batchMatch(const std::vector<std::vector<int32_t>>& tokens) const;
+
+  void reset() {
+    std::unique_lock<std::mutex> lock(mutex_);
+    if (trie_) {
+      trie_->reset();
+    }
+  }
+
+  const Param& param() const {
+    return param_;
+  }
+
+ private:
+  void insertWorker();
+};
+
+}  // namespace ngram

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

@@ -0,0 +1,104 @@
+#pragma once
+
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+
+#include <tvm/ffi/reflection/registry.h>
+
+#include "ngram.h"
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <stdexcept>
+#include <vector>
+
+struct NgramCorpusObj : public tvm::ffi::Object {
+ public:
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.NgramCorpus", NgramCorpusObj, tvm::ffi::Object);
+  static constexpr bool _type_mutable = true;
+
+  NgramCorpusObj(
+      int64_t capacity,
+      int64_t max_trie_depth,
+      int64_t min_bfs_breadth,
+      int64_t max_bfs_breadth,
+      int64_t draft_token_num,
+      int64_t match_type) {
+    ngram::Param param;
+    param.enable = true;
+    param.enable_router_mode = false;
+    param.max_trie_depth = static_cast<size_t>(max_trie_depth);
+    param.min_bfs_breadth = static_cast<size_t>(min_bfs_breadth);
+    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";
+    ngram_ = std::make_unique<ngram::Ngram>(static_cast<size_t>(capacity), param);
+  }
+
+  void async_insert(const tvm::ffi::TensorView tokens_flat, const tvm::ffi::TensorView offsets) {
+    auto* data = static_cast<const int32_t*>(tokens_flat.data_ptr());
+    auto* offs = static_cast<const int64_t*>(offsets.data_ptr());
+    int64_t batch_size = offsets.size(0) - 1;
+
+    std::vector<std::vector<int32_t>> tokens(batch_size);
+    for (int64_t i = 0; i < batch_size; ++i) {
+      tokens[i].assign(data + offs[i], data + offs[i + 1]);
+    }
+    ngram_->asyncInsert(std::move(tokens));
+  }
+
+  void batch_match(
+      const tvm::ffi::TensorView tokens_flat,
+      const tvm::ffi::TensorView offsets,
+      const tvm::ffi::TensorView out_tokens,
+      const tvm::ffi::TensorView out_mask) {
+    auto* data = static_cast<const int32_t*>(tokens_flat.data_ptr());
+    auto* offs = static_cast<const int64_t*>(offsets.data_ptr());
+    int64_t batch_size = offsets.size(0) - 1;
+
+    std::vector<std::vector<int32_t>> tokens(batch_size);
+    for (int64_t i = 0; i < batch_size; ++i) {
+      tokens[i].assign(data + offs[i], data + offs[i + 1]);
+    }
+
+    auto result = ngram_->batchMatch(tokens);
+
+    auto* out_tok = static_cast<int32_t*>(out_tokens.data_ptr());
+    auto* out_msk = static_cast<uint8_t*>(out_mask.data_ptr());
+    if (result.token.size() > static_cast<size_t>(out_tokens.size(0))) {
+      throw std::runtime_error(
+          "out_tokens buffer too small: " + std::to_string(out_tokens.size(0)) + " < " +
+          std::to_string(result.token.size()));
+    }
+    if (result.mask.size() > static_cast<size_t>(out_mask.size(0))) {
+      throw std::runtime_error(
+          "out_mask buffer too small: " + std::to_string(out_mask.size(0)) + " < " +
+          std::to_string(result.mask.size()));
+    }
+    std::memcpy(out_tok, result.token.data(), result.token.size() * sizeof(int32_t));
+    std::memcpy(out_msk, result.mask.data(), result.mask.size() * sizeof(uint8_t));
+  }
+
+  void synchronize() {
+    ngram_->synchronize();
+  }
+
+  void reset() {
+    ngram_->reset();
+  }
+
+ private:
+  std::unique_ptr<ngram::Ngram> ngram_;
+};
+
+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("async_insert", &NgramCorpusObj::async_insert)
+      .def("batch_match", &NgramCorpusObj::batch_match)
+      .def("synchronize", &NgramCorpusObj::synchronize)
+      .def("reset", &NgramCorpusObj::reset);
+}
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(register_once, register_ngram_corpus);

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

@@ -0,0 +1,104 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdlib>
+#include <iostream>
+#include <limits>
+#include <regex>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace ngram {
+
+struct Param {
+  bool enable;
+  bool enable_router_mode;
+  size_t min_bfs_breadth;
+  size_t max_bfs_breadth;
+  size_t max_trie_depth;
+  size_t draft_token_num;
+  std::string match_type;
+
+  std::vector<size_t> batch_draft_token_num;
+
+  size_t get_draft_token_num(size_t batch_size) const {
+    if (batch_size < batch_draft_token_num.size()) {
+      if (batch_draft_token_num[batch_size] !=
+          std::numeric_limits<decltype(batch_draft_token_num)::value_type>::max()) {
+        return batch_draft_token_num[batch_size];
+      }
+    }
+    return draft_token_num - 1;
+  }
+
+  std::vector<size_t> parse(const std::string& value) {
+    // 0-1|10,2-3|20,
+    std::vector<size_t> result;
+    if (value.empty()) {
+      return result;
+    }
+    std::vector<size_t> mark;
+    std::regex comma_re(",");
+    std::sregex_token_iterator first{value.begin(), value.end(), comma_re, -1}, last;
+    for (auto p : std::vector<std::string>(first, last)) {
+      std::cerr << "seg " << p << std::endl;
+    }
+    for (const auto& seg : std::vector<std::string>(first, last)) {
+      std::regex pipe_re("\\|");
+      std::sregex_token_iterator seg_first{seg.begin(), seg.end(), pipe_re, -1}, seg_last;
+      std::vector<std::string> part(seg_first, seg_last);
+      for (auto p : part) {
+        std::cerr << "part " << p << std::endl;
+      }
+      if (part.size() != 2) {
+        throw std::runtime_error(
+            "failed to get config, invalid config: " + seg + ", part's size = " + std::to_string(part.size()));
+      }
+      std::regex endash_re("-");
+      std::sregex_token_iterator range_first{part[0].begin(), part[0].end(), endash_re, -1}, range_last;
+      std::vector<std::string> range(range_first, range_last);
+      if (range.size() != 2) {
+        throw std::runtime_error("failed to get range, invalid config: " + value);
+      }
+      size_t L = std::atoi(range[0].c_str());
+      size_t R = std::atoi(range[1].c_str());
+      if (L > R || R > 128) {
+        throw std::runtime_error("invalid range, config: " + value);
+      }
+      if (R >= result.size()) {
+        result.resize(R + 1, std::numeric_limits<decltype(result)::value_type>::max());
+        mark.resize(result.size(), false);
+      }
+      size_t config = std::atoi(part[1].c_str());
+      do {
+        if (mark[L]) {
+          throw std::runtime_error("repeated position " + std::to_string(L) + ", config : " + value);
+        }
+        mark[L] = true;
+        result[L] = config;
+      } while (++L <= R);
+    }
+    return result;
+  }
+
+  void resetBatchReturnTokenNum(const std::string& value) {
+    batch_draft_token_num = parse(value);
+  }
+
+  std::string detail() {
+    std::stringstream ss;
+    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;
+    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] << ",";
+    }
+    return ss.str();
+  }
+};
+
+}  // namespace ngram

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

@@ -0,0 +1,73 @@
+#pragma once
+
+#include <condition_variable>
+#include <mutex>
+#include <queue>
+#include <utility>
+
+namespace utils {
+
+template <typename T>
+class Queue {
+ public:
+  bool enqueue(T&& rhs) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      if (closed_) {
+        return false;
+      }
+      queue_.emplace(std::move(rhs));
+    }
+    cv_.notify_one();
+    return true;
+  }
+
+  bool enqueue(const T& rhs) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      if (closed_) {
+        return false;
+      }
+      queue_.emplace(rhs);
+    }
+    cv_.notify_one();
+    return true;
+  }
+
+  bool dequeue(T& rhs) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    cv_.wait(lock, [this] { return queue_.size() || closed_; });
+    if (closed_) {
+      return false;
+    }
+    rhs = std::move(queue_.front());
+    queue_.pop();
+    return true;
+  }
+
+  size_t size() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return queue_.size();
+  }
+
+  bool empty() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return queue_.empty();
+  }
+
+  void close() {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      closed_ = true;
+    }
+    cv_.notify_all();
+  }
+
+ private:
+  std::queue<T> queue_;
+  mutable std::mutex mutex_;
+  std::condition_variable cv_;
+  bool closed_{false};
+};
+
+}  // namespace utils

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

@@ -0,0 +1,61 @@
+#include "result.h"
+
+#include <cstring>
+#include <queue>
+#include <tuple>
+
+namespace ngram {
+
+Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root) {
+  Result info;
+  std::vector<int32_t> prevs;
+  info.token.reserve(draft_token_num);
+  prevs.reserve(draft_token_num);
+  std::queue<std::tuple<int32_t, int32_t, int32_t>> queue;
+  info.token.emplace_back(last_token);
+  prevs.emplace_back(-1);
+
+  for (auto [token, next] : tree[root].next) {
+    queue.emplace(token, next, 0);
+  }
+  while (queue.size()) {
+    auto [token, next, prev] = queue.front();
+    queue.pop();
+    info.token.emplace_back(token);
+    prevs.emplace_back(prev);
+    for (auto [t, n] : tree[next].next) {
+      queue.emplace(t, n, info.token.size() - 1);
+    }
+  }
+
+  // zero padding to length
+  while (info.token.size() < static_cast<size_t>(draft_token_num)) {
+    info.token.emplace_back(0);
+    prevs.emplace_back(0);
+  }
+
+  int n = info.token.size();
+  info.mask.resize(n * n, 0);
+  info.mask[0] = 1;
+  for (int i = 0; i < n; ++i) {
+    if (prevs[i] != -1) {
+      memcpy(&info.mask[i * n], &info.mask[prevs[i] * n], prevs[i] + 1);
+    }
+    info.mask[i * n + i] = 1;
+  }
+
+  return info;
+}
+
+void Result::truncate(size_t n) {
+  if (n < token.size()) {
+    int full_n = token.size();
+    for (size_t i = 1; i < n; ++i) {
+      memcpy(&mask[i * n], &mask[i * full_n], sizeof(mask[0]) * n);
+    }
+    token.resize(n);
+    mask.resize(n * n);
+  }
+}
+
+}  // namespace ngram

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

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <cstdint>
+#include <unordered_map>
+#include <vector>
+
+namespace ngram {
+
+struct Result {
+  std::vector<int32_t> token;
+  std::vector<uint8_t> mask;
+
+  void truncate(size_t n);
+};
+
+struct Node {
+  std::unordered_map<int32_t, int32_t> next;
+};
+
+Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root);
+
+}  // namespace ngram

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

@@ -0,0 +1,230 @@
+#include "trie.h"
+
+#include <algorithm>
+#include <cstring>
+#include <list>
+#include <queue>
+#include <tuple>
+#include <vector>
+
+namespace ngram {
+
+Trie::Trie(size_t capacity, const Param& param) : param_(param) {
+  nodes_.resize(capacity);
+  for (auto& node : nodes_) {
+    node_pool_.emplace_back(&node);
+  }
+  free_node_count_ = node_pool_.size();
+  root_ = getNode();
+}
+
+void Trie::insert(const int32_t* tokens, size_t len) {
+  for (size_t i = 0; i < len; ++i) {
+    auto start = tokens + i;
+    auto end = start + std::min(len - i, param_.max_trie_depth);
+
+    if (static_cast<size_t>(end - start) > free_node_count_) {
+      squeeze(end - start - free_node_count_);
+    }
+
+    TrieNode* cursor = root_;
+    path_.clear();
+    while (start != end) {
+      auto token = *start;
+      auto iter = cursor->child.find(token);
+      if (iter == cursor->child.end()) {
+        iter = cursor->child.insert({token, getNode()}).first;
+        auto node = iter->second;
+
+        cursor->lru.emplace_front(node);
+        global_lru_.emplace_back(node);
+
+        node->token = token;
+        node->parent = cursor;
+        node->parent_lru_pos = cursor->lru.begin();
+        node->global_lru_pos = --global_lru_.end();
+        node->freq = 1;
+        cursor->sorted_children.insert(node);
+      } else {
+        auto node = iter->second;
+        cursor->sorted_children.erase(node);
+        node->freq++;
+        cursor->sorted_children.insert(node);
+        cursor->lru.splice(cursor->lru.begin(), cursor->lru, node->parent_lru_pos);
+      }
+      cursor = iter->second;
+      path_.emplace_back(cursor);
+      ++start;
+    }
+
+    for (auto it = path_.rbegin(); it != path_.rend(); ++it) {
+      TrieNode* node = *it;
+      global_lru_.splice(global_lru_.begin(), global_lru_, node->global_lru_pos);
+    }
+  }
+}
+
+void Trie::squeeze(size_t count) {
+  if (!(node_pool_.size() >= free_node_count_ + count)) {
+    throw std::runtime_error(
+        "Insufficient node size to release required nodes. "
+        "available to release: " +
+        std::to_string(node_pool_.size() - free_node_count_) + ", required to release: " + std::to_string(count));
+  }
+  while (count--) {
+    auto last = global_lru_.back();
+    global_lru_.pop_back();
+
+    if (!last->child.empty()) {
+      throw std::runtime_error(
+          "The node to be released still has child nodes and cannot be "
+          "released. ");
+    }
+
+    last->parent->lru.erase(last->parent_lru_pos);
+    last->parent->sorted_children.erase(last);
+    last->parent->child.erase(last->token);
+
+    node_pool_[free_node_count_++] = last;
+  }
+}
+
+void Trie::reset() {
+  global_lru_.clear();
+  path_.clear();
+  node_pool_.clear();
+  for (auto& node : nodes_) {
+    node_pool_.emplace_back(&node);
+  }
+  free_node_count_ = node_pool_.size();
+  root_ = getNode();
+}
+
+std::vector<std::pair<TrieNode*, int32_t>> Trie::match(const int32_t* context, size_t len) const {
+  std::vector<std::pair<TrieNode*, int32_t>> result;
+  const auto max_match_depth = std::min(len, param_.max_trie_depth);
+  result.reserve(max_match_depth);
+  for (size_t match_depth = max_match_depth; match_depth > 0; --match_depth) {
+    auto start = context + len - match_depth;
+    auto end = start + match_depth;
+    auto cursor = root_;
+    while (start != end) {
+      auto iter = cursor->child.find(*start);
+      if (iter == cursor->child.end()) {
+        cursor = nullptr;
+        break;
+      }
+      ++start;
+      cursor = iter->second;
+    }
+    if (cursor != nullptr && !cursor->child.empty()) {
+      result.emplace_back(cursor, static_cast<int32_t>(match_depth));
+    }
+  }
+  return result;
+}
+
+Result Trie::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);
+
+  const auto max_match_depth = std::max<int32_t>(1, static_cast<int32_t>(param.max_trie_depth - 1));
+  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 (auto [node, depth] : anchors) {
+    std::queue<std::tuple<int32_t, double, const TrieNode*>> queue;
+    queue.push({root, (max_match_depth - depth) * bfs_breadth_scale + param.min_bfs_breadth, node});
+    while (queue.size() && cursor <= static_cast<int>(draft_token_num)) {
+      auto front = queue.front();
+      queue.pop();
+
+      auto parent = std::get<0>(front);
+      auto cur_breadth = std::get<1>(front);
+      auto iter = std::get<2>(front)->lru.begin();
+
+      auto breadth = std::max(1, int32_t(cur_breadth));
+      for (int i = 0;
+           i < breadth && iter != std::get<2>(front)->lru.end() && cursor <= static_cast<int>(draft_token_num);
+           ++i, ++iter) {
+        auto token = (*iter)->token;
+        auto pos = -1;
+        if (auto tit = tree[parent].next.find(token); tit != tree[parent].next.end()) {
+          pos = tit->second;
+        } else {
+          pos = tree[parent].next.insert(std::make_pair(token, cursor++)).first->second;
+        }
+        queue.emplace(pos, cur_breadth - bfs_breadth_scale, *iter);
+      }
+    }
+  }
+
+  return fillResult(last_token, draft_token_num + 1, tree, root);
+}
+
+Result Trie::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);
+
+  struct CompareByLastDouble {
+    bool operator()(
+        const std::tuple<double, const TrieNode*, double>& a,
+        const std::tuple<double, const TrieNode*, double>& b) const {
+      return std::get<2>(a) < std::get<2>(b);
+    }
+  };
+
+  std::priority_queue<
+      std::tuple<double, const TrieNode*, double>,
+      std::vector<std::tuple<double, const TrieNode*, double>>,
+      CompareByLastDouble>
+      heap;
+
+  std::vector<Node> tree(draft_token_num + 1);
+
+  int root = 0;
+  int cursor = 1;
+  int top_k = param.max_bfs_breadth;
+
+  auto addToHeap = [&heap, &top_k](int parent, const TrieNode* trie_node, double prob) -> void {
+    double sum_freq = 0.0;
+    int count = 0;
+    std::list<std::pair<TrieNode*, int32_t>> topk_children;
+    for (auto* child : trie_node->sorted_children) {
+      sum_freq += static_cast<double>(child->freq);
+      topk_children.emplace_back(child, child->freq);
+      if (++count >= top_k) break;
+    }
+    if (sum_freq <= 0) sum_freq = 1.0;
+    for (const auto& [child, freq] : topk_children) {
+      double norm_freq = static_cast<double>(freq) / sum_freq * prob;
+      heap.emplace(parent, child, norm_freq);
+    }
+  };
+
+  for (auto [node, _] : anchors) {
+    addToHeap(root, node, 1.0);
+
+    while (!heap.empty() && cursor <= static_cast<int>(draft_token_num)) {
+      auto [parent, trie_node, prob] = heap.top();
+      heap.pop();
+      auto token = trie_node->token;
+      int pos = -1;
+      auto tit = tree[parent].next.find(token);
+      if (tit != tree[parent].next.end()) {
+        pos = tit->second;
+      } else {
+        pos = cursor++;
+        tree[parent].next[token] = pos;
+      }
+      addToHeap(pos, trie_node, prob);
+    }
+  }
+
+  return fillResult(last_token, draft_token_num + 1, tree, root);
+}
+
+}  // namespace ngram

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

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "param.h"
+#include "result.h"
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <list>
+#include <new>
+#include <set>
+#include <tuple>
+#include <unordered_map>
+#include <vector>
+
+namespace ngram {
+
+struct TrieNode {
+  std::unordered_map<int32_t, TrieNode*> child;
+  std::list<TrieNode*>::const_iterator global_lru_pos;
+  std::list<TrieNode*>::const_iterator parent_lru_pos;
+  int32_t token;
+  TrieNode* parent;
+  std::list<TrieNode*> lru;
+  int32_t freq = 0;
+
+  struct CompareByFreq {
+    bool operator()(TrieNode* a, TrieNode* b) const {
+      return std::tie(b->freq, a->token, a) < std::tie(a->freq, b->token, b);
+    }
+  };
+  std::multiset<TrieNode*, CompareByFreq> sorted_children;
+};
+
+class Trie {
+ public:
+  Trie(size_t capacity, const Param& param);
+
+  void insert(const int32_t* tokens, size_t len);
+
+  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;
+
+  void squeeze(size_t count);
+
+  void reset();
+
+ private:
+  std::vector<std::pair<TrieNode*, int32_t>> match(const int32_t* context, size_t len) const;
+
+  TrieNode* getNode() {
+    auto node = node_pool_[--free_node_count_];
+    node->~TrieNode();
+    new (node) TrieNode();
+    return node;
+  }
+
+  std::vector<TrieNode> nodes_;
+  std::vector<TrieNode*> node_pool_;
+  size_t free_node_count_;
+  std::list<TrieNode*> global_lru_;
+  TrieNode* root_;
+  std::vector<TrieNode*> path_;
+  Param param_;
+};
+
+}  // namespace ngram