sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/speculative/cpp_ngram/ngram.h

@@ -0,0 +1,110 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <list>
+#include <mutex>
+#include <set>
+#include <sstream>
+#include <thread>
+#include <tuple>
+#include <unordered_map>
+#include <vector>
+
+#include "param.h"
+#include "queue.h"
+
+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 Ngram {
+  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_;
+
+  std::vector<std::pair<TrieNode*, int32_t>> match(const std::vector<int32_t>& tokens, size_t batch_size) const;
+
+  void squeeze(size_t count);
+
+  TrieNode* getNode() {
+    auto node = node_pool_[--free_node_count_];
+    node->~TrieNode();
+    new (node) TrieNode();
+    return node;
+  }
+
+  mutable std::mutex mutex_;
+  bool quit_flag_;
+  utils::Queue<std::vector<int32_t>> insert_queue_;
+  std::thread insert_worker_;
+  std::vector<std::tuple<int32_t, int32_t, int32_t, int32_t>> match_tmp_data_;
+
+ public:
+  Ngram(size_t capacity, const Param& param);
+  Ngram() = default;
+  ~Ngram();
+
+  static Ngram& instance() {
+    static Ngram instance;
+    return instance;
+  }
+
+  void synchronize() const;
+
+  void asyncInsert(std::vector<std::vector<int32_t>>&& tokens);
+
+  struct Result {
+    std::vector<int32_t> token;
+    std::vector<uint8_t> mask;
+
+    void truncate(size_t n);
+  };
+
+  Result batchMatch(const std::vector<std::vector<int32_t>>& tokens) const;
+
+  void reset() {
+    std::unique_lock<std::mutex> lock(mutex_);
+
+    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();
+  }
+
+  const Param& param() const {
+    return param_;
+  }
+
+ private:
+  Result matchBFS(const std::vector<int32_t>& tokens, size_t batch_size) const;
+  Result matchProb(const std::vector<int32_t>& tokens, size_t batch_size) const;
+
+  void insert();
+};
+
+}  // namespace ngram

sglang/srt/speculative/cpp_ngram/ngram_cache.py

@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+
+import logging
+import os
+from typing import List, Tuple
+
+import numpy as np
+from torch.utils.cpp_extension import load
+
+logger = logging.getLogger(__name__)
+
+_abs_path = os.path.dirname(os.path.abspath(__file__))
+ngram_cache_cpp = load(
+    name="ngram_cache_cpp",
+    sources=[
+        f"{_abs_path}/ngram_cache_binding.cpp",
+        f"{_abs_path}/ngram.cpp",
+    ],
+    extra_cflags=["-O3", "-std=c++20"],
+)
+
+
+class NgramCache:
+    def __init__(
+        self,
+        branch_length=18,
+        min_match_window_size=1,
+        max_match_window_size=10,
+        min_bfs_breadth=1,
+        max_bfs_breadth=8,
+        draft_token_num=8,
+        match_type="BFS",
+        capacity=1000000,
+    ):
+        param = ngram_cache_cpp.Param()
+        param.branch_length = branch_length
+        param.min_match_window_size = min_match_window_size
+        param.max_match_window_size = max_match_window_size
+        param.min_bfs_breadth = min_bfs_breadth
+        param.max_bfs_breadth = max_bfs_breadth
+        param.draft_token_num = draft_token_num
+        param.match_type = match_type
+        self.cache = ngram_cache_cpp.Ngram(capacity, param)
+
+        self.default_mask = np.ones((1, 1), dtype=np.int64)
+        self.draft_token_num = draft_token_num
+
+    def batch_put(self, batch_tokens: List[List[int]]):
+        self.cache.asyncInsert(batch_tokens)
+
+    def synchronize(self):
+        self.cache.synchronize()
+
+    def reset(self):
+        self.cache.reset()
+
+    def batch_get(self, batch_tokens: List[List[int]]) -> Tuple[np.ndarray, np.ndarray]:
+        result = self.cache.batchMatch(batch_tokens)
+        return np.array(result.token), np.array(result.mask)
+
+    def leaf_paths_from_mask(
+        self, tokens: List[int], tree_mask: List[List[int]]
+    ) -> List[List[int]]:
+        """
+        Find all leaf paths according to the binary tree_mask (i.e., paths that are not prefixes of any other path).
+
+        Args:
+            mask   : List[List[int]]   # nxn binary matrix
+            tokens : List[int]         # token list corresponding to columns
+
+        Returns:
+            List[List[int]]            # token lists of only the leaf paths, preserving their order of appearance
+        """
+
+        row_sets = [
+            (i, {idx for idx, v in enumerate(row) if v == 1})
+            for i, row in enumerate(tree_mask)
+        ]
+        leaf_sets = []
+        leaf_rows = []
+
+        for i, cur_set in reversed(row_sets):
+            if any(cur_set <= kept for kept in leaf_sets):
+                continue
+            leaf_sets.append(cur_set)
+            leaf_rows.append(i)
+
+        leaf_rows.reverse()
+        result = []
+        for r in leaf_rows:
+            path = [tokens[col] for col in range(len(tokens)) if tree_mask[r][col] == 1]
+            result.append(path)
+
+        return result
+
+    def debug_result(
+        self, decoding_ids: np.ndarray, decoding_masks: np.ndarray, tokenizer=None
+    ):
+        decoding_ids = decoding_ids.reshape(-1, self.draft_token_num)
+        decoding_masks = decoding_masks.reshape(
+            -1, self.draft_token_num, self.draft_token_num
+        )
+        logger.info(f"\n{decoding_ids=}\n{decoding_masks=}")
+        for i in range(decoding_ids.shape[0]):
+            leaf_paths = self.leaf_paths_from_mask(
+                decoding_ids[i].tolist(), decoding_masks[i].tolist()
+            )
+            if tokenizer is None:
+                logger.info(f"draft path {i}: {leaf_paths}")
+            else:
+                logger.info(f"result {i}:")
+                for leaf_path in leaf_paths:
+                    logger.info(
+                        f"draft path {i}: {leaf_path} -> {tokenizer.decode(leaf_path, ensure_ascii=False)}"
+                    )
+
+
+# main function
+if __name__ == "__main__":
+    format = f"%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s"
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format=format,
+        datefmt="%Y-%m-%d %H:%M:%S",
+        force=True,
+    )
+
+    token_ids = [
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+        [1, 2, 3, 44, 55, 66, 77, 88, 99, 100],
+    ]
+    cache = NgramCache(branch_length=12, draft_token_num=8)
+    cache.batch_put(token_ids)
+
+    cache.synchronize()
+    decoding_ids, decoding_masks = cache.batch_get([[1, 2, 3], [3, 44], [3, 6, 999]])
+
+    cache.debug_result(decoding_ids, decoding_masks)

sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp

@@ -0,0 +1,43 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "ngram.h"
+
+PYBIND11_MODULE(ngram_cache_cpp, m) {
+  using namespace ngram;
+  namespace py = pybind11;
+  m.doc() = "";
+
+  py::class_<Ngram>(m, "Ngram")
+      .def(py::init<size_t, const Param&>(), py::arg("capacity"), py::arg("param"))
+      .def("asyncInsert", &Ngram::asyncInsert, "")
+      .def("batchMatch", &Ngram::batchMatch, "")
+      .def("reset", &Ngram::reset, "")
+      .def("synchronize", &Ngram::synchronize, "");
+
+  py::class_<Param>(m, "Param")
+      .def(py::init<>())
+      .def_readwrite("enable", &Param::enable)
+      .def_readwrite("enable_router_mode", &Param::enable_router_mode)
+      .def_readwrite("min_bfs_breadth", &Param::min_bfs_breadth)
+      .def_readwrite("max_bfs_breadth", &Param::max_bfs_breadth)
+      .def_readwrite("min_match_window_size", &Param::min_match_window_size)
+      .def_readwrite("max_match_window_size", &Param::max_match_window_size)
+      .def_readwrite("branch_length", &Param::branch_length)
+      .def_readwrite("draft_token_num", &Param::draft_token_num)
+      .def_readwrite("match_type", &Param::match_type)
+      .def_readwrite("batch_min_match_window_size", &Param::batch_min_match_window_size)
+      .def_readwrite("batch_draft_token_num", &Param::batch_draft_token_num)
+      .def("get_draft_token_num", &Param::get_draft_token_num, "")
+      .def("get_min_match_window_size", &Param::get_min_match_window_size, "")
+      .def("parse", &Param::parse, "")
+      .def("resetBatchMinMatchWindowSize", &Param::resetBatchMinMatchWindowSize, "")
+      .def("resetBatchReturnTokenNum", &Param::resetBatchReturnTokenNum, "")
+      .def("detail", &Param::detail, "");
+
+  py::class_<Ngram::Result>(m, "Result")
+      .def(py::init<>())
+      .def_readwrite("token", &Ngram::Result::token)
+      .def_readwrite("mask", &Ngram::Result::mask)
+      .def("truncate", &Ngram::Result::truncate);
+}

sglang/srt/speculative/cpp_ngram/param.h

@@ -0,0 +1,125 @@
+#pragma once
+
+#include <cstddef>
+#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 min_match_window_size;
+  size_t max_match_window_size;
+  size_t branch_length;
+  size_t draft_token_num;
+  std::string match_type;
+
+  std::vector<size_t> batch_min_match_window_size;
+  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;
+  }
+
+  size_t get_min_match_window_size(size_t batch_size) const {
+    if (batch_size < batch_min_match_window_size.size()) {
+      if (batch_min_match_window_size[batch_size] !=
+          std::numeric_limits<decltype(batch_min_match_window_size)::value_type>::max()) {
+        return batch_min_match_window_size[batch_size];
+      }
+    }
+    return min_match_window_size;
+  }
+
+  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 resetBatchMinMatchWindowSize(const std::string& value) {
+    batch_min_match_window_size = parse(value);
+  }
+
+  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
+       << ", min_match_window_size = " << min_match_window_size << ", max_match_window_size = " << max_match_window_size
+       << ", branch_length = " << branch_length << ", draft_token_num = " << draft_token_num
+       << ", match_type = " << match_type;
+    ss << ", batch_min_match_window_size(" << batch_min_match_window_size.size() << ") = ";
+    for (int i = 0; i < batch_min_match_window_size.size(); ++i) {
+      ss << i << "|" << batch_min_match_window_size[i] << ",";
+    }
+    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

sglang/srt/speculative/cpp_ngram/queue.h

@@ -0,0 +1,71 @@
+#pragma once
+
+#include <condition_variable>
+#include <queue>
+
+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

sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -20,7 +20,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
 )
-from sglang.srt.speculative.eagle_utils import EagleDraftInput
+from sglang.srt.speculative.eagle_info import EagleDraftInput
 from sglang.srt.utils import (
     require_attn_tp_gather,
     require_gathered_buffer,
@@ -91,6 +91,9 @@ class EAGLEDraftCudaGraphRunner:
                 (self.max_num_token * self.speculative_num_steps,), dtype=torch.int64
             )
             self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            self.mrope_positions = torch.zeros(
+                (3, self.max_num_token), dtype=torch.int64
+            )
             self.topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
             self.topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64)
             self.hidden_states = torch.zeros(
@@ -159,6 +162,7 @@ class EAGLEDraftCudaGraphRunner:
         seq_lens = self.seq_lens[:num_seqs]
         out_cache_loc = self.out_cache_loc[: num_tokens * self.speculative_num_steps]
         positions = self.positions[:num_tokens]
+        mrope_positions = self.mrope_positions[:, :num_tokens]
         topk_p = self.topk_p[:num_seqs]
         topk_index = self.topk_index[:num_seqs]
         hidden_states = self.hidden_states[:num_seqs]
@@ -224,6 +228,7 @@ class EAGLEDraftCudaGraphRunner:
             seq_lens_sum=seq_lens.sum().item(),
             return_logprob=False,
             positions=positions,
+            mrope_positions=mrope_positions,
             global_num_tokens_gpu=global_num_tokens,
             dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
             global_dp_buffer_len=global_dp_buffer_len,
@@ -297,6 +302,7 @@ class EAGLEDraftCudaGraphRunner:
         if bs != raw_bs:
             self.seq_lens.fill_(self.seq_len_fill_value)
             self.out_cache_loc.zero_()
+            self.positions.zero_()
 
         num_tokens = bs * self.num_tokens_per_bs
 

sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -21,7 +21,8 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
 )
-from sglang.srt.speculative.eagle_utils import EagleDraftInput, fast_topk
+from sglang.srt.speculative.eagle_info import EagleDraftInput
+from sglang.srt.speculative.spec_utils import fast_topk
 from sglang.srt.utils import (
     require_attn_tp_gather,
     require_gathered_buffer,
@@ -80,6 +81,9 @@ class EAGLEDraftExtendCudaGraphRunner:
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
             self.out_cache_loc = torch.ones((self.max_num_token,), dtype=torch.int64)
             self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            self.mrope_positions = torch.zeros(
+                (3, self.max_num_token), dtype=torch.int64
+            )
 
             if self.eagle_worker.speculative_algorithm.is_eagle3():
                 self.hidden_states = torch.zeros(
@@ -189,6 +193,7 @@ class EAGLEDraftExtendCudaGraphRunner:
         accept_length = self.accept_length[:bs]
         out_cache_loc = self.out_cache_loc[:num_tokens]
         positions = self.positions[:num_tokens]
+        mrope_positions = self.mrope_positions[:, :num_tokens]
         hidden_states = self.hidden_states[:num_tokens]
         next_token_logits_buffer = self.next_token_logits_buffer[:bs]
 
@@ -247,6 +252,7 @@ class EAGLEDraftExtendCudaGraphRunner:
             seq_lens_sum=seq_lens.sum().item(),
             return_logprob=False,
             positions=positions,
+            mrope_positions=mrope_positions,
             global_num_tokens_gpu=self.global_num_tokens_gpu,
             global_num_tokens_for_logprob_gpu=self.global_num_tokens_for_logprob_gpu,
             dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
@@ -326,6 +332,7 @@ class EAGLEDraftExtendCudaGraphRunner:
         if bs * self.num_tokens_per_bs != num_tokens:
             self.seq_lens.fill_(self.seq_len_fill_value)
             self.out_cache_loc.zero_()
+            self.positions.zero_()
             self.accept_length.fill_(1)
             self.extend_seq_lens.fill_(1)
 
@@ -336,7 +343,11 @@ class EAGLEDraftExtendCudaGraphRunner:
             self.extend_seq_lens[:raw_bs].copy_(forward_batch.extend_seq_lens)
         self.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
         self.positions[:num_tokens].copy_(forward_batch.positions)
-        self.hidden_states[:num_tokens].copy_(forward_batch.spec_info.hidden_states)
+        if (
+            forward_batch.spec_info.hidden_states.shape[1]
+            == self.hidden_states.shape[1]
+        ):
+            self.hidden_states[:num_tokens].copy_(forward_batch.spec_info.hidden_states)
         if forward_batch.spec_info.accept_length is not None:
             self.accept_length[:raw_bs].copy_(forward_batch.spec_info.accept_length)
         self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)