perf: batch draft traversal and compact grammar caches

cpp/compiled_grammar.cc

@@ -185,7 +185,9 @@ std::optional<SerializationError> DeserializeJSONValue(
   if (object.find("grammar") == object.end()) {
     return ConstructDeserializeError("Expect a 'grammar' field", type_name);
   }
-  AutoDeserializeJSONValue(&(impl->grammar), object["grammar"], type_name);
+  if (auto error = AutoDeserializeJSONValue(&(impl->grammar), object["grammar"], type_name)) {
+    return error;
+  }
   if (object.find("tokenizer_metadata") == object.end()) {
     return ConstructDeserializeError("Expect a 'tokenizer_metadata' field", type_name);
   }
@@ -199,7 +201,11 @@ std::optional<SerializationError> DeserializeJSONValue(
   if (object.find("adaptive_token_mask_cache") == object.end()) {
     return ConstructDeserializeError("Expect a 'adaptive_token_mask_cache' field", type_name);
   }
-  AutoDeserializeJSONValue(&(impl->adaptive_token_mask_cache), object["adaptive_token_mask_cache"]);
+  if (auto error = AutoDeserializeJSONValue(
+          &(impl->adaptive_token_mask_cache), object["adaptive_token_mask_cache"], type_name
+      )) {
+    return error;
+  }
   return std::nullopt;
 }
 

cpp/fsm.h

@@ -618,6 +618,9 @@ class FSMWithStartEndBase {
    */
   const std::vector<int32_t>& GetEnds() const { return ends_; }
 
+  /*! \brief Returns whether this FSM view is deterministic. */
+  bool GetIsDFA() const { return is_dfa_; }
+
   /*!
    * \brief Checks if a given state is an end/accepting state.
    * \param state The state to check.

cpp/grammar.cc

@@ -5,7 +5,12 @@
 
 #include <xgrammar/grammar.h>
 
+#include <algorithm>
+#include <cstdint>
+#include <optional>
 #include <string>
+#include <utility>
+#include <vector>
 
 #include "grammar_functor.h"
 #include "grammar_parser.h"
@@ -20,6 +25,32 @@
 
 namespace xgrammar {
 
+struct CompactRuleFSMView {
+  int start = 0;
+  std::vector<int32_t> ends;
+  bool is_dfa = false;
+  int edge_num = 0;
+  int node_num = 0;
+
+  CompactRuleFSMView() = default;
+
+  explicit CompactRuleFSMView(const CompactFSMWithStartEndWithSize& value)
+      : start(value.GetFsm().GetStart()),
+        ends(value.GetFsm().GetEnds()),
+        is_dfa(value.GetFsm().GetIsDFA()),
+        edge_num(value.GetEdgeNum()),
+        node_num(value.GetNodeNum()) {}
+};
+
+XGRAMMAR_MEMBER_ARRAY(
+    CompactRuleFSMView,
+    &CompactRuleFSMView::start,
+    &CompactRuleFSMView::ends,
+    &CompactRuleFSMView::is_dfa,
+    &CompactRuleFSMView::edge_num,
+    &CompactRuleFSMView::node_num
+);
+
 /******************* Grammar::Impl *******************/
 
 std::size_t MemorySize(const Grammar::Impl& impl) {
@@ -33,6 +64,132 @@ std::size_t MemorySize(const Grammar::Impl& impl) {
          MemorySize(impl.allow_empty_rule_ids);
 }
 
+picojson::value SerializeJSONValue(const Grammar::Impl& impl) {
+  std::vector<std::optional<CompactRuleFSMView>> rule_fsm_views;
+  rule_fsm_views.reserve(impl.per_rule_fsms.size());
+  for (const auto& rule_fsm : impl.per_rule_fsms) {
+    if (rule_fsm.has_value()) {
+      rule_fsm_views.emplace_back(CompactRuleFSMView(*rule_fsm));
+    } else {
+      rule_fsm_views.emplace_back(std::nullopt);
+    }
+  }
+
+  picojson::object result;
+  result["rules"] = AutoSerializeJSONValue(impl.rules_);
+  result["suffix_stop_infos"] = AutoSerializeJSONValue(impl.suffix_stop_infos_);
+  // Preserve the historical public field labels. The internal vector names
+  // are inverted relative to those labels.
+  result["grammar_expr_data"] = AutoSerializeJSONValue(impl.grammar_expr_indptr_);
+  result["grammar_expr_indptr"] = AutoSerializeJSONValue(impl.grammar_expr_data_);
+  result["root_rule_id"] = AutoSerializeJSONValue(impl.root_rule_id_);
+  result["complete_fsm"] = AutoSerializeJSONValue(impl.complete_fsm);
+  result["per_rule_fsms"] = AutoSerializeJSONValue(rule_fsm_views);
+  result["allow_empty_rule_ids"] = AutoSerializeJSONValue(impl.allow_empty_rule_ids);
+  result["optimized"] = AutoSerializeJSONValue(impl.optimized);
+  return picojson::value(std::move(result));
+}
+
+std::optional<SerializationError> DeserializeJSONValue(
+    Grammar::Impl* impl, const picojson::value& value, const std::string& type_name
+) {
+  if (!value.is<picojson::object>()) {
+    return ConstructDeserializeError("Expect an object", type_name);
+  }
+  const auto& object = value.get<picojson::object>();
+
+  auto deserialize_field = [&](const char* name,
+                               auto* destination) -> std::optional<SerializationError> {
+    const auto it = object.find(name);
+    if (it == object.end()) {
+      return ConstructDeserializeError("Missing member " + std::string(name), type_name);
+    }
+    return AutoDeserializeJSONValue(destination, it->second, type_name);
+  };
+
+  if (auto error = deserialize_field("rules", &impl->rules_)) {
+    return error;
+  }
+  if (auto error = deserialize_field("suffix_stop_infos", &impl->suffix_stop_infos_)) {
+    return error;
+  }
+  if (auto error = deserialize_field("grammar_expr_data", &impl->grammar_expr_indptr_)) {
+    return error;
+  }
+  if (auto error = deserialize_field("grammar_expr_indptr", &impl->grammar_expr_data_)) {
+    return error;
+  }
+  if (auto error = deserialize_field("root_rule_id", &impl->root_rule_id_)) {
+    return error;
+  }
+  if (auto error = deserialize_field("complete_fsm", &impl->complete_fsm)) {
+    return error;
+  }
+  if (auto error = deserialize_field("allow_empty_rule_ids", &impl->allow_empty_rule_ids)) {
+    return error;
+  }
+  if (auto error = deserialize_field("optimized", &impl->optimized)) {
+    return error;
+  }
+
+  std::vector<std::optional<CompactRuleFSMView>> rule_fsm_views;
+  if (auto error = deserialize_field("per_rule_fsms", &rule_fsm_views)) {
+    return error;
+  }
+  const std::size_t expected_rule_fsm_count = impl->optimized ? impl->rules_.size() : 0;
+  if (rule_fsm_views.size() != expected_rule_fsm_count) {
+    return ConstructDeserializeError(
+        "per_rule_fsms count does not match grammar optimization state", type_name
+    );
+  }
+  if (!rule_fsm_views.empty() && impl->complete_fsm.IsNull()) {
+    return ConstructDeserializeError("optimized grammar is missing the complete FSM", type_name);
+  }
+
+  const int num_states = rule_fsm_views.empty() ? 0 : impl->complete_fsm.NumStates();
+  const std::size_t num_edges = rule_fsm_views.empty() ? 0 : impl->complete_fsm.GetNumEdges();
+  impl->per_rule_fsms.clear();
+  impl->per_rule_fsms.reserve(rule_fsm_views.size());
+  for (auto& rule_fsm : rule_fsm_views) {
+    if (!rule_fsm.has_value()) {
+      impl->per_rule_fsms.emplace_back(std::nullopt);
+      continue;
+    }
+    if (rule_fsm->start < 0 || rule_fsm->start >= num_states) {
+      return ConstructDeserializeError("per-rule FSM start state is out of range", type_name);
+    }
+    if (rule_fsm->edge_num < 0 || rule_fsm->node_num < 0) {
+      return ConstructDeserializeError("per-rule FSM size is negative", type_name);
+    }
+    if (static_cast<std::size_t>(rule_fsm->edge_num) > num_edges ||
+        rule_fsm->node_num > num_states) {
+      return ConstructDeserializeError("per-rule FSM size exceeds the complete FSM", type_name);
+    }
+    if (!std::is_sorted(rule_fsm->ends.begin(), rule_fsm->ends.end()) ||
+        std::adjacent_find(rule_fsm->ends.begin(), rule_fsm->ends.end()) != rule_fsm->ends.end()) {
+      return ConstructDeserializeError(
+          "per-rule FSM end states are not sorted and unique", type_name
+      );
+    }
+    for (const int32_t end : rule_fsm->ends) {
+      if (end < 0 || end >= num_states) {
+        return ConstructDeserializeError("per-rule FSM end state is out of range", type_name);
+      }
+    }
+
+    CompactFSMWithStartEnd fsm(
+        impl->complete_fsm, rule_fsm->start, std::move(rule_fsm->ends), rule_fsm->is_dfa
+    );
+    impl->per_rule_fsms.emplace_back(
+        CompactFSMWithStartEndWithSize(std::move(fsm), rule_fsm->edge_num, rule_fsm->node_num)
+    );
+  }
+  impl->per_rule_fsm_hashes.clear();
+  impl->per_rule_fsm_new_state_ids.clear();
+
+  return std::nullopt;
+}
+
 /******************* Grammar *******************/
 
 std::string Grammar::ToString() const { return GrammarPrinter(*this).ToString(); }

cpp/grammar_compiler.cc

@@ -10,6 +10,8 @@
 #include <cctype>
 #include <cstddef>
 #include <cstdint>
+#include <future>
+#include <memory>
 #include <optional>
 #include <string>
 #include <unordered_map>
@@ -1016,8 +1018,8 @@ class GrammarCompilerSub {
       std::optional<RuleLevelCache> rule_level_cache
   )
       : tokenizer_info_(tokenizer_info),
-        max_threads_(max_threads),
-        rule_level_cache_(rule_level_cache) {}
+        rule_level_cache_(rule_level_cache),
+        thread_pool_(max_threads > 1 ? std::make_unique<ThreadPool>(max_threads) : nullptr) {}
 
   CompiledGrammar CompileBuiltinJSONGrammar();
 
@@ -1054,11 +1056,12 @@ class GrammarCompilerSub {
 
   /*! \brief The vocabulary associated with this storage class. */
   const TokenizerInfo tokenizer_info_;
-  /*! \brief The maximum number of threads to use. */
-  const int max_threads_;
 
   /*! \brief The manager of the rule level cache.*/
   std::optional<RuleLevelCache> rule_level_cache_;
+
+  /*! \brief A persistent pool shared by concurrent grammar compilations. */
+  std::unique_ptr<ThreadPool> thread_pool_;
 };
 
 CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_unoptimized) {
@@ -1081,13 +1084,33 @@ CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_un
   // 2. All byte strings (with element_in_string=0, 1, 2, ...)
   // since other positions will be expanded to the above positions
 
-  // TODO(Charlie): Figure out how to support ThreadPool and std::mutex in WebAssembly.
-  // Only declare ThreadPool and mutex if max_threads > 1, so when max_threads = 1, we do
-  // not need ThreadPool or std::mutex, which throws error in runtime in WebAssembly.
-  std::optional<ThreadPool> thread_pool;
+  auto root_rule_id = compiled_grammar_impl->grammar->GetRootRuleId();
+  std::vector<std::pair<ParserState, bool>> adaptive_states;
+  for (int32_t rule_id = 0; rule_id < static_cast<int>(compiled_grammar_impl->grammar->NumRules());
+       ++rule_id) {
+    auto rule = compiled_grammar_impl->grammar->GetRule(rule_id);
+    const auto& rule_fsm = compiled_grammar_impl->grammar->per_rule_fsms[rule_id];
+    XGRAMMAR_DCHECK(rule_fsm.has_value());
+    auto cur_stack_element =
+        ParserState(rule_id, rule.body_expr_id, 0, ParserState::kNoPrevInputPos, 0);
+    std::unordered_set<int> reachable_states;
+    rule_fsm->GetFsm().GetReachableStates(&reachable_states);
+    for (int state : reachable_states) {
+      cur_stack_element.element_id = state;
+      if (!rule_fsm->GetFsm().IsScanableState(state)) {
+        continue;
+      }
+      adaptive_states.emplace_back(cur_stack_element, rule_id == root_rule_id);
+    }
+  }
+
+  // Small grammars compile faster inline on independent outer compiler workers. Larger
+  // FSMs amortize native queueing and use the shared pool, which bounds total native
+  // concurrency across simultaneous compilations.
+  constexpr std::size_t kMinParallelStates = 64;
+  const bool use_thread_pool = thread_pool_ && adaptive_states.size() >= kMinParallelStates;
   std::optional<std::mutex> adaptive_token_mask_cache_mutex;
-  if (max_threads_ > 1) {
-    thread_pool.emplace(max_threads_);
+  if (use_thread_pool) {
     adaptive_token_mask_cache_mutex.emplace();
   }
 
@@ -1100,47 +1123,45 @@ CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_un
         rule_level_cache_
     );
     auto cur_adaptive_token_mask_cache = grammar_matcher.GetAdaptiveTokenMask(is_root_rule);
-    if (max_threads_ > 1) {
+    if (use_thread_pool) {
       std::lock_guard<std::mutex> lock(adaptive_token_mask_cache_mutex.value());
       compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache;
     } else {
       compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache;
     }
   };
 
-  auto add_task_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) {
-    // Execute depending on whether we use thread_pool
-    if (max_threads_ > 1) {
-      thread_pool->Execute([add_adaptive_token_mask, state, is_root_rule]() {
-        add_adaptive_token_mask(state, is_root_rule);
-      });
-    } else {
+  if (!use_thread_pool) {
+    for (const auto& [state, is_root_rule] : adaptive_states) {
       add_adaptive_token_mask(state, is_root_rule);
     }
-  };
-
-  auto root_rule_id = compiled_grammar_impl->grammar->GetRootRuleId();
-
-  for (int32_t rule_id = 0; rule_id < static_cast<int>(compiled_grammar_impl->grammar->NumRules());
-       ++rule_id) {
-    auto rule = compiled_grammar_impl->grammar->GetRule(rule_id);
-    const auto& rule_fsm = compiled_grammar_impl->grammar->per_rule_fsms[rule_id];
-    XGRAMMAR_DCHECK(rule_fsm.has_value());
-    auto cur_stack_element =
-        ParserState(rule_id, rule.body_expr_id, 0, ParserState::kNoPrevInputPos, 0);
-    std::unordered_set<int> reachable_states;
-    rule_fsm->GetFsm().GetReachableStates(&reachable_states);
-    for (int i : reachable_states) {
-      cur_stack_element.element_id = i;
-      if (!rule_fsm->GetFsm().IsScanableState(i)) {
-        continue;
-      }
-      add_task_adaptive_token_mask(cur_stack_element, rule_id == root_rule_id);
-    }
+    return CompiledGrammar(compiled_grammar_impl);
   }
 
-  if (max_threads_ > 1) {
-    thread_pool->Join();
+  std::vector<std::shared_future<void>> pending_tasks;
+  pending_tasks.reserve(adaptive_states.size());
+  try {
+    for (const auto& adaptive_state : adaptive_states) {
+      // Capturing names introduced by a structured binding requires C++20.
+      // XGrammar targets C++17, so copy the pair members into ordinary locals.
+      const ParserState state = adaptive_state.first;
+      const bool is_root_rule = adaptive_state.second;
+      pending_tasks.emplace_back(
+          thread_pool_->Submit([add_adaptive_token_mask, state, is_root_rule]() {
+            add_adaptive_token_mask(state, is_root_rule);
+          })
+      );
+    }
+    for (auto& task : pending_tasks) {
+      task.get();
+    }
+  } catch (...) {
+    // Tasks capture function-local compilation state. Drain the shared pool before
+    // unwinding so no task can outlive that state.
+    if (thread_pool_) {
+      thread_pool_->Wait();
+    }
+    throw;
   }
 
   return CompiledGrammar(compiled_grammar_impl);