[Feature] Provide batched fill next token mask and batched accept string/token methods.

cpp/grammar_matcher.cc

@@ -8,7 +8,10 @@
 
 #include <xgrammar/matcher.h>
 
+#include <algorithm>
 #include <cstdint>
+#include <optional>
+#include <thread>
 #include <utility>
 #include <vector>
 
@@ -19,6 +22,7 @@
 #include "support/encoding.h"
 #include "support/int_set.h"
 #include "support/logging.h"
+#include "support/thread_pool.h"
 #include "testing.h"
 
 namespace xgrammar {
@@ -292,6 +296,26 @@ class GrammarMatcher::Impl : public EarleyParser {
 
   std::string _DebugPrintInternalState() const { return PrintStates(); }
 
+  static void BatchFillNextTokenBitmask(
+      std::vector<GrammarMatcher>* matchers,
+      DLTensor* next_token_bitmask,
+      const std::optional<std::vector<int32_t>>& indices = std::nullopt,
+      int max_thread = 16,
+      bool debug_print = false
+  );
+
+  static std::vector<uint8_t> BatchedAcceptString(
+      std::vector<GrammarMatcher>* matchers,
+      const std::vector<std::string>& input_strs,
+      bool debug_print = false
+  );
+
+  static std::vector<uint8_t> BatchedAcceptToken(
+      std::vector<GrammarMatcher>* matchers,
+      const std::vector<int32_t>& token_ids,
+      bool debug_print = false
+  );
+
  private:
   using StoreType = AdaptiveTokenMask::StoreType;
 
@@ -852,6 +876,76 @@ int GrammarMatcher::Impl::GetNextUncertainToken(
   }
 }
 
+void GrammarMatcher::Impl::BatchFillNextTokenBitmask(
+    std::vector<GrammarMatcher>* matchers,
+    DLTensor* next_token_bitmask,
+    const std::optional<std::vector<int32_t>>& indices,
+    int max_thread,
+    bool debug_print
+) {
+  XGRAMMAR_CHECK(!indices.has_value() || indices->size() == matchers->size())
+      << "The size of indices (" << (indices.has_value() ? indices->size() : 0)
+      << ") should be the same as the size of matchers (" << matchers->size() << ").";
+
+  if (max_thread == 1) {
+    for (int i = 0; i < static_cast<int32_t>(matchers->size()); i++) {
+      auto& matcher = (*matchers)[i];
+      int index = indices.has_value() ? (*indices)[i] : i;
+      XGRAMMAR_CHECK(index >= 0 && index < next_token_bitmask->shape[0])
+          << "The index " << index << " is out of range [0, " << next_token_bitmask->shape[0]
+          << ") for batch_id " << i << ".";
+      matcher->FillNextTokenBitmask(next_token_bitmask, index, debug_print);
+    }
+  } else {
+    XGRAMMAR_CHECK(max_thread > 0);
+    ThreadPool thread_pool(
+        std::min(std::thread::hardware_concurrency(), static_cast<uint32_t>(max_thread))
+    );
+    auto fill_next_token_mask = [&](int32_t batch_id) {
+      auto& matcher = (*matchers)[batch_id];
+      int index = indices.has_value() ? (*indices)[batch_id] : batch_id;
+      XGRAMMAR_CHECK(index >= 0 && index < next_token_bitmask->shape[0])
+          << "The index " << index << " is out of range [0, " << next_token_bitmask->shape[0]
+          << ") for batch_id " << batch_id << ".";
+      matcher->FillNextTokenBitmask(next_token_bitmask, index, debug_print);
+    };
+    for (int i = 0; i < static_cast<int32_t>(matchers->size()); i++) {
+      thread_pool.Execute([fill_next_token_mask, i]() { fill_next_token_mask(i); });
+    }
+    thread_pool.Join();
+  }
+}
+
+std::vector<uint8_t> GrammarMatcher::Impl::BatchedAcceptString(
+    std::vector<GrammarMatcher>* matchers,
+    const std::vector<std::string>& input_strs,
+    bool debug_print
+) {
+  XGRAMMAR_CHECK(matchers->size() == input_strs.size())
+      << "The size of matchers (" << matchers->size() << ") and input_strs (" << input_strs.size()
+      << ") should be the same.";
+  std::vector<uint8_t> accepted(matchers->size());
+  for (int i = 0; i < static_cast<int32_t>(matchers->size()); i++) {
+    auto& matcher = (*matchers)[i];
+    accepted[i] = matcher->AcceptString(input_strs[i], debug_print);
+  }
+  return accepted;
+}
+
+std::vector<uint8_t> GrammarMatcher::Impl::BatchedAcceptToken(
+    std::vector<GrammarMatcher>* matchers, const std::vector<int32_t>& token_ids, bool debug_print
+) {
+  XGRAMMAR_CHECK(matchers->size() == token_ids.size())
+      << "The size of matchers (" << matchers->size() << ") and token_ids (" << token_ids.size()
+      << ") should be the same.";
+  std::vector<uint8_t> accepted(matchers->size());
+  for (int i = 0; i < static_cast<int32_t>(matchers->size()); i++) {
+    auto& matcher = (*matchers)[i];
+    accepted[i] = matcher->AcceptToken(token_ids[i], debug_print);
+  }
+  return accepted;
+}
+
 GrammarMatcher::GrammarMatcher(
     const CompiledGrammar& compiled_grammar,
     std::optional<std::vector<int>> override_stop_tokens,
@@ -894,4 +988,28 @@ std::string GrammarMatcher::_DebugPrintInternalState() const {
   return pimpl_->_DebugPrintInternalState();
 }
 
+void GrammarMatcher::BatchFillNextTokenBitmask(
+    std::vector<GrammarMatcher>* matchers,
+    DLTensor* next_token_bitmask,
+    const std::optional<std::vector<int32_t>>& indices,
+    int max_thread,
+    bool debug_print
+) {
+  Impl::BatchFillNextTokenBitmask(matchers, next_token_bitmask, indices, max_thread, debug_print);
+}
+
+std::vector<uint8_t> GrammarMatcher::BatchAcceptString(
+    std::vector<GrammarMatcher>* matchers,
+    const std::vector<std::string>& input_strs,
+    bool debug_print
+) {
+  return Impl::BatchedAcceptString(matchers, input_strs, debug_print);
+}
+
+std::vector<uint8_t> GrammarMatcher::BatchAcceptToken(
+    std::vector<GrammarMatcher>* matchers, const std::vector<int32_t>& token_ids, bool debug_print
+) {
+  return Impl::BatchedAcceptToken(matchers, token_ids, debug_print);
+}
+
 }  // namespace xgrammar

cpp/nanobind/nanobind.cc

@@ -13,12 +13,19 @@
 #include <nanobind/stl/vector.h>
 #include <xgrammar/xgrammar.h>
 
+#include <cstdint>
+#include <optional>
+#include <thread>
+#include <variant>
+#include <vector>
+
 #include "../grammar_functor.h"
 #include "../json_schema_converter.h"
 #include "../regex_converter.h"
 #include "../testing.h"
 #include "python_methods.h"
 #include "xgrammar/exception.h"
+#include "xgrammar/matcher.h"
 
 namespace nb = nanobind;
 
@@ -74,6 +81,67 @@ bool GrammarMatcher_FillNextTokenBitmask(
   return matcher.FillNextTokenBitmask(bitmask_dltensor_ptr, index, debug_print);
 }
 
+void GrammarMatcher_BatchFillNextTokenMask(
+    std::vector<GrammarMatcher>* matchers,
+    nb::ndarray<> arr,
+    const std::optional<std::vector<int32_t>>& indices,
+    std::variant<int32_t, std::string> max_thread,
+    bool debug_print
+) {
+  if (arr.ndim() != 2) {
+    throw std::runtime_error("batch_token_bitmask tensor must be 2D");
+  }
+  if (arr.device_type() != nb::device::cpu::value) {
+    throw std::runtime_error("token_bitmask array must be on CPU");
+  }
+  if (arr.dtype() != nb::dtype<int32_t>()) {
+    throw std::runtime_error("token_bitmask array must be int32");
+  }
+  static_assert(sizeof(arr) == sizeof(void*) + sizeof(nb::dlpack::dltensor));
+
+  DLTensor* bitmask_dltensor_ptr =
+      reinterpret_cast<::DLTensor*>(reinterpret_cast<char*>(&arr) + sizeof(void*));
+
+  int unwrapped_max_thread;
+
+  if (std::holds_alternative<std::string>(max_thread)) {
+    if (std::get<std::string>(max_thread) == "auto") {
+      unwrapped_max_thread = std::min(std::thread::hardware_concurrency() / 2, 1u);
+    } else {
+      throw std::runtime_error("max_thread must be an integer or 'auto'");
+    }
+  } else {
+    XGRAMMAR_CHECK(std::get<int32_t>(max_thread) > 0) << "max_thread must be positive";
+    unwrapped_max_thread = std::get<int32_t>(max_thread);
+  }
+
+  GrammarMatcher::BatchFillNextTokenBitmask(
+      matchers, bitmask_dltensor_ptr, indices, unwrapped_max_thread, debug_print
+  );
+}
+
+std::vector<uint8_t> GrammarMatcher_BatchAcceptString(
+    std::vector<GrammarMatcher>* matchers,
+    const std::vector<std::variant<nb::bytes, std::string>>& input_strs,
+    bool debug_print
+) {
+  if (matchers->size() != input_strs.size()) {
+    throw std::runtime_error("The size of matchers and input_strs must be the same");
+  }
+  std::vector<uint8_t> accepted(matchers->size());
+  for (int i = 0; i < static_cast<int32_t>(matchers->size()); i++) {
+    std::string input_str;
+    if (std::holds_alternative<std::string>(input_strs[i])) {
+      input_str = std::get<std::string>(input_strs[i]);
+    } else {
+      input_str = std::get<nb::bytes>(input_strs[i]).c_str();
+    }
+    auto& matcher = (*matchers)[i];
+    accepted[i] = matcher.AcceptString(input_str, debug_print);
+  }
+  return accepted;
+}
+
 std::vector<nanobind::bytes> TokenizerInfo_GetDecodedVocab(const TokenizerInfo& tokenizer) {
   const auto& decoded_vocab = tokenizer.GetDecodedVocab();
   std::vector<nanobind::bytes> py_result;
@@ -233,6 +301,26 @@ NB_MODULE(xgrammar_bindings, m) {
           &GrammarMatcher_FillNextTokenBitmask,
           nb::call_guard<nb::gil_scoped_release>()
       )
+      .def_static(
+          "batch_fill_next_token_bitmask",
+          &GrammarMatcher_BatchFillNextTokenMask,
+          nb::arg("matchers"),
+          nb::arg("batch_token_bitmask"),
+          nb::arg("indices").none(),
+          nb::arg("max_thread") = "auto",
+          nb::arg("debug_print") = false,
+          nb::call_guard<nb::gil_scoped_release>()
+      )
+      .def_static(
+          "batch_accept_string",
+          &GrammarMatcher_BatchAcceptString,
+          nb::call_guard<nb::gil_scoped_release>()
+      )
+      .def_static(
+          "batch_accept_token",
+          &GrammarMatcher::BatchAcceptToken,
+          nb::call_guard<nb::gil_scoped_release>()
+      )
       .def(
           "find_jump_forward_string",
           &GrammarMatcher::FindJumpForwardString,

docs/xgrammar_features/runtime_safeguards.md

@@ -15,7 +15,7 @@ If the recursion depth exceeds the limit,
 the matcher operations (including [`xgr.GrammarMatcher.accept_token`](xgrammar.GrammarMatcher.accept_token),
 [`xgr.GrammarMatcher.accept_string`](xgrammar.GrammarMatcher.accept_string), [`xgr.GrammarMatcher.fill_next_token_bitmask`](xgrammar.GrammarMatcher.fill_next_token_bitmask),
 [`xgr.GrammarMatcher.find_jump_forward_string`](xgrammar.GrammarMatcher.find_jump_forward_string)) will raise
-`RuntimeError`.
+`RuntimeError`(before XGrammar v0.1.21).
 
 You can also use the [`xgr.max_recursion_depth`](xgrammar.max_recursion_depth) context manager to set the maximum
 recursion depth for a code block.
@@ -27,6 +27,8 @@ with max_recursion_depth(10000):
     matcher.accept_token(token_id)
 ```
 
+After XGrammar v0.1.21, the pushdown automaton parser was replaced with an Earley parser, and the recursion depth cannot be infinite anymore. Thus, the feature was removed.
+
 ## Cache Size Limit
 
 The {py:class}`xgr.GrammarCompiler` class uses a cache to store the compiled grammars.

include/xgrammar/matcher.h

@@ -142,6 +142,50 @@ class GrammarMatcher {
    */
   std::string _DebugPrintInternalState() const;
 
+  /*!
+    \brief A batched version of FillNextTokenBitmask for better efficiency.
+    \param matchers The array of GrammarMatcher objects.
+    \param next_token_bitmask The pre-allocated DLTensor to store the result bitmasks.
+    \param indices The optional array of indices to specify which matcher corresponds to which slice
+    of the bitmask tensor. If not provided, all matchers will write to the corresponding
+    indices(matchers[i] to next_token_bitmask[i]).
+    \param max_thread The maximum number of threads to use. Default is 16.
+    \param debug_print Whether to print debug information. Default is false.
+  */
+  static void BatchFillNextTokenBitmask(
+      std::vector<GrammarMatcher>* matchers,
+      DLTensor* next_token_bitmask,
+      const std::optional<std::vector<int32_t>>& indices = std::nullopt,
+      int max_thread = 16,
+      bool debug_print = false
+  );
+
+  /*!
+   * \brief A batched version of AcceptString for better efficiency.
+   * \param matchers The array of GrammarMatcher objects.
+   * \param input_strs The array of input strings to be accepted.
+   * \param debug_print Whether to print debug information. Default is false.
+   * \return A vector of bytes indicating whether each string is accepted.
+   */
+  static std::vector<uint8_t> BatchAcceptString(
+      std::vector<GrammarMatcher>* matchers,
+      const std::vector<std::string>& input_strs,
+      bool debug_print = false
+  );
+
+  /*!
+   * \brief A batched version of AcceptToken for better efficiency.
+   * \param matchers The array of GrammarMatcher objects.
+   * \param token_ids The array of token ids to be accepted.
+   * \param debug_print Whether to print debug information. Default is false.
+   * \return A vector of bytes indicating whether each token is accepted.
+   */
+  static std::vector<uint8_t> BatchAcceptToken(
+      std::vector<GrammarMatcher>* matchers,
+      const std::vector<int32_t>& token_ids,
+      bool debug_print = false
+  );
+
   XGRAMMAR_DEFINE_PIMPL_METHODS(GrammarMatcher);
 };