[Fix] Fix the optimization for TagDispatch and SubGrammarAdder. (#471)

This PR fixes a concurrency problem in the optimization for TagDispatch,
and fixes a bug in `SubGrammarAdder`.

---------

Signed-off-by: Yuchuan <[email protected]>

Author: Seven-Streams

cpp/grammar_compiler.cc

@@ -636,13 +636,20 @@ class GrammarCompilerNoCache {
  private:
   /*! \brief The main logic. Compile the grammar with multi-threading. */
   CompiledGrammar MultiThreadCompileGrammar(Grammar grammar);
+  /*! \brief Optimization for TagDispatch.
+   *  \param compiled_grammar_impl the compiled_grammar to be optimized.
+   *  \param tag_dispatch_rule_id_to_second_slicing_bitset Return value. Mapping from the rule_id to
+   * the definite accepted token mask.
+   */
+  void TagDispatchOptimization(
+      std::shared_ptr<CompiledGrammar::Impl> compiled_grammar_impl,
+      std::unordered_map<int32_t, DynamicBitset>* tag_dispatch_rule_id_to_second_slicing_bitset
+  );
 
   /*! \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 Mapping from the rule_id to the definite accepted token mask. */
-  std::unordered_map<int32_t, DynamicBitset> tag_dispatch_rule_id_to_second_slicing_bitset;
 };
 
 CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar grammar_unoptimized) {
@@ -655,48 +662,8 @@ CompiledGrammar GrammarCompilerNoCache::MultiThreadCompileGrammar(Grammar gramma
   if (tokenizer_info_.GetVocabSize() == 0) {
     return CompiledGrammar(compiled_grammar_impl);
   }
-
-  // Optimization for TagDispatch: Precompute the definitely accepted tokens.
-  for (int i = 0; i < compiled_grammar_impl->grammar->NumRules(); i++) {
-    const auto& rule = compiled_grammar_impl->grammar->GetRule(i);
-    const auto& rule_body = compiled_grammar_impl->grammar->GetGrammarExpr(rule.body_expr_id);
-    if (rule_body.type != GrammarExprType::kTagDispatch) {
-      continue;
-    }
-    XGRAMMAR_DCHECK(rule_body.type == GrammarExprType::kTagDispatch);
-    Grammar::Impl::TagDispatch tag_dispatch =
-        compiled_grammar_impl->grammar->GetTagDispatch(rule.body_expr_id);
-    const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab();
-    DynamicBitset definite_accepted_tokens_since_second_char(sorted_decoded_vocab.size());
-    for (int i = 0; i < static_cast<int32_t>(sorted_decoded_vocab.size()); i++) {
-      bool definite_accept_since_second_char = true;
-      const auto& token = sorted_decoded_vocab[i].second;
-      if (token.empty()) {
-        definite_accepted_tokens_since_second_char.Set(i);
-        continue;
-      }
-
-      // Check if the token contains any tag or stop string after the first character.
-      for (const auto& tag : tag_dispatch.tag_rule_pairs) {
-        if (token.find(tag.first, 1) != std::string::npos) {
-          definite_accept_since_second_char = false;
-          break;
-        }
-      }
-      for (const auto& stop_str : tag_dispatch.stop_str) {
-        if (token.find(stop_str, 1) != std::string::npos) {
-          definite_accept_since_second_char = false;
-          break;
-        }
-      }
-
-      // If the token can be definitely accepted since the second character, set the bit.
-      if (definite_accept_since_second_char) {
-        definite_accepted_tokens_since_second_char.Set(i);
-      }
-    }
-    tag_dispatch_rule_id_to_second_slicing_bitset[i] = definite_accepted_tokens_since_second_char;
-  }
+  std::unordered_map<int32_t, DynamicBitset> tag_dispatch_rule_id_to_second_slicing_bitset;
+  TagDispatchOptimization(compiled_grammar_impl, &tag_dispatch_rule_id_to_second_slicing_bitset);
   // Step 3. Compute the adaptive token mask cache
   // The token mask cache is computed for these positions in the grammar:
   // 1. All character class or character class star (with last_utf8_bytes=0, 1, 2, 3)
@@ -842,6 +809,57 @@ CompiledGrammar GrammarCompilerNoCache::CompileGrammar(
   return MultiThreadCompileGrammar(Grammar::FromEBNF(ebnf_str, root_rule_name));
 }
 
+void GrammarCompilerNoCache::TagDispatchOptimization(
+    std::shared_ptr<CompiledGrammar::Impl> compiled_grammar_impl,
+    std::unordered_map<int32_t, DynamicBitset>* tag_dispatch_rule_id_to_second_slicing_bitset
+) {
+  using GrammarExprType = Grammar::Impl::GrammarExprType;
+  tag_dispatch_rule_id_to_second_slicing_bitset->clear();
+
+  // Optimization for TagDispatch: Precompute the definitely accepted tokens.
+  for (int i = 0; i < compiled_grammar_impl->grammar->NumRules(); i++) {
+    const auto& rule = compiled_grammar_impl->grammar->GetRule(i);
+    const auto& rule_body = compiled_grammar_impl->grammar->GetGrammarExpr(rule.body_expr_id);
+    if (rule_body.type != GrammarExprType::kTagDispatch) {
+      continue;
+    }
+    XGRAMMAR_DCHECK(rule_body.type == GrammarExprType::kTagDispatch);
+    Grammar::Impl::TagDispatch tag_dispatch =
+        compiled_grammar_impl->GetGrammar()->GetTagDispatch(rule.body_expr_id);
+    const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab();
+    DynamicBitset definite_accepted_tokens_since_second_char(sorted_decoded_vocab.size());
+    for (int i = 0; i < static_cast<int32_t>(sorted_decoded_vocab.size()); i++) {
+      bool definite_accept_since_second_char = true;
+      const auto& token = sorted_decoded_vocab[i].second;
+      if (token.empty()) {
+        definite_accepted_tokens_since_second_char.Set(i);
+        continue;
+      }
+
+      // Check if the token contains any tag or stop string after the first character.
+      for (const auto& tag : tag_dispatch.tag_rule_pairs) {
+        if (token.find(tag.first, 1) != std::string::npos) {
+          definite_accept_since_second_char = false;
+          break;
+        }
+      }
+      for (const auto& stop_str : tag_dispatch.stop_str) {
+        if (token.find(stop_str, 1) != std::string::npos) {
+          definite_accept_since_second_char = false;
+          break;
+        }
+      }
+
+      // If the token can be definitely accepted since the second character, set the bit.
+      if (definite_accept_since_second_char) {
+        definite_accepted_tokens_since_second_char.Set(i);
+      }
+    }
+    (*tag_dispatch_rule_id_to_second_slicing_bitset)[i] =
+        definite_accepted_tokens_since_second_char;
+  }
+}
+
 /******************* GrammarCompiler::Impl *******************/
 
 /*!

cpp/grammar_functor.cc

@@ -73,6 +73,18 @@ class SubGrammarAdderImpl : public GrammarMutator {
     );
   }
 
+  int32_t VisitTagDispatch(const GrammarExpr& grammar_expr) final {
+    Grammar::Impl::TagDispatch old_tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr);
+    Grammar::Impl::TagDispatch new_tag_dispatch;
+    new_tag_dispatch.stop_eos = old_tag_dispatch.stop_eos;
+    for (const auto& [tag, rule_id] : old_tag_dispatch.tag_rule_pairs) {
+      new_tag_dispatch.tag_rule_pairs.emplace_back(tag, new_rule_ids_names[rule_id].first);
+    }
+    new_tag_dispatch.stop_str = old_tag_dispatch.stop_str;
+    new_tag_dispatch.loop_after_dispatch = old_tag_dispatch.loop_after_dispatch;
+    return builder_->AddTagDispatch(new_tag_dispatch);
+  }
+
   std::vector<std::pair<int32_t, std::string>> new_rule_ids_names;
 };
 

tests/python/test_grammar_matcher_structural_tag.py

@@ -1,5 +1,6 @@
 import json
 import sys
+import threading
 import time
 from typing import List
 
@@ -340,5 +341,36 @@ def test_utf8_structural_tag_begin_end():
     _ = compiler.compile_structural_tag(structures, triggers)
 
 
+@pytest.mark.hf_token_required
+def test_pressure_structural_tag():
+    model = "meta-llama/Llama-3.1-8B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True, trust_remote_code=True)
+    tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer)
+    compiler = xgr.GrammarCompiler(tokenizer_info, max_threads=1)
+    threads = []
+    start = "start"
+    schema = {"type": "object", "properties": {"arg": {"type": "string"}}}
+    end = "end"
+
+    def worker(idx: int):
+        tag = xgr.StructuralTagItem(begin=start, schema=schema, end=end)
+        triggers = [start]
+        stag_grammar = xgr.Grammar.from_structural_tag([tag], triggers)
+        start_grammar = xgr.Grammar.from_ebnf("root ::= [a-z] root | [a-z]")
+        grammar = start_grammar
+        for _ in range(idx):
+            grammar = grammar.concat(grammar, start_grammar)
+        final_grammar = xgr.Grammar.concat(grammar, stag_grammar)
+        _ = compiler.compile_grammar(final_grammar)
+
+    for i in range(128):
+        t = threading.Thread(target=worker, args=(i,))
+        threads.append(t)
+        t.start()
+
+    for t in threads:
+        t.join()
+
+
 if __name__ == "__main__":
     pytest.main(sys.argv)

tests/python/test_grammar_union_concat.py

@@ -83,5 +83,84 @@ def test_grammar_concat():
     assert str(concat_grammar) == expected
 
 
+def test_grammar_union_with_stag():
+    expected_grammar_union = r"""root ::= ((root_1_1) | (root_2))
+basic_escape ::= (([\"\\/bfnrt]) | ("u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9]))
+basic_string_sub ::= (("\"") | ([^\0-\x1f\"\\\r\n] basic_string_sub) | ("\\" basic_escape basic_string_sub)) (=([ \n\t]* [,}\]:]))
+basic_any ::= ((basic_number) | (basic_string) | (basic_boolean) | (basic_null) | (basic_array) | (basic_object))
+basic_integer ::= (("0") | (basic_integer_1 [1-9] [0-9]*))
+basic_number ::= ((basic_number_1 basic_number_7 basic_number_3 basic_number_6))
+basic_string ::= (("\"" basic_string_sub))
+basic_boolean ::= (("true") | ("false"))
+basic_null ::= (("null"))
+basic_array ::= (("[" [ \n\t]* basic_any basic_array_1 [ \n\t]* "]") | ("[" [ \n\t]* "]"))
+basic_object ::= (("{" [ \n\t]* basic_string [ \n\t]* ":" [ \n\t]* basic_any basic_object_1 [ \n\t]* "}") | ("{" [ \n\t]* "}"))
+root_1 ::= (("{" [ \n\t]* "\"arg\"" [ \n\t]* ":" [ \n\t]* basic_string [ \n\t]* "}") | ("{" [ \n\t]* "}"))
+basic_integer_1 ::= ("" | ("-"))
+basic_number_1 ::= ("" | ("-"))
+basic_number_2 ::= (([0-9] basic_number_2) | ([0-9]))
+basic_number_3 ::= ("" | ("." basic_number_2))
+basic_number_4 ::= ("" | ([+\-]))
+basic_number_5 ::= (([0-9] basic_number_5) | ([0-9]))
+basic_number_6 ::= ("" | ([eE] basic_number_4 basic_number_5))
+basic_array_1 ::= ("" | ([ \n\t]* "," [ \n\t]* basic_any basic_array_1))
+basic_object_1 ::= ("" | ([ \n\t]* "," [ \n\t]* basic_string [ \n\t]* ":" [ \n\t]* basic_any basic_object_1))
+basic_number_7 ::= (("0") | ([1-9] [0-9]*))
+triggered_tags_group ::= (("" root_1 "end"))
+triggered_tags ::= TagDispatch(
+  ("start", triggered_tags_group),
+  stop_eos=true,
+  stop_str=(),
+  loop_after_dispatch=true
+)
+root_1_1 ::= ((triggered_tags))
+root_2 ::= (([a-z] root_2) | ([a-z]))
+"""
+
+    expected_grammar_concat = r"""root ::= ((root_1_1 root_2))
+basic_escape ::= (([\"\\/bfnrt]) | ("u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9]))
+basic_string_sub ::= (("\"") | ([^\0-\x1f\"\\\r\n] basic_string_sub) | ("\\" basic_escape basic_string_sub)) (=([ \n\t]* [,}\]:]))
+basic_any ::= ((basic_number) | (basic_string) | (basic_boolean) | (basic_null) | (basic_array) | (basic_object))
+basic_integer ::= (("0") | (basic_integer_1 [1-9] [0-9]*))
+basic_number ::= ((basic_number_1 basic_number_7 basic_number_3 basic_number_6))
+basic_string ::= (("\"" basic_string_sub))
+basic_boolean ::= (("true") | ("false"))
+basic_null ::= (("null"))
+basic_array ::= (("[" [ \n\t]* basic_any basic_array_1 [ \n\t]* "]") | ("[" [ \n\t]* "]"))
+basic_object ::= (("{" [ \n\t]* basic_string [ \n\t]* ":" [ \n\t]* basic_any basic_object_1 [ \n\t]* "}") | ("{" [ \n\t]* "}"))
+root_1 ::= (("{" [ \n\t]* "\"arg\"" [ \n\t]* ":" [ \n\t]* basic_string [ \n\t]* "}") | ("{" [ \n\t]* "}"))
+basic_integer_1 ::= ("" | ("-"))
+basic_number_1 ::= ("" | ("-"))
+basic_number_2 ::= (([0-9] basic_number_2) | ([0-9]))
+basic_number_3 ::= ("" | ("." basic_number_2))
+basic_number_4 ::= ("" | ([+\-]))
+basic_number_5 ::= (([0-9] basic_number_5) | ([0-9]))
+basic_number_6 ::= ("" | ([eE] basic_number_4 basic_number_5))
+basic_array_1 ::= ("" | ([ \n\t]* "," [ \n\t]* basic_any basic_array_1))
+basic_object_1 ::= ("" | ([ \n\t]* "," [ \n\t]* basic_string [ \n\t]* ":" [ \n\t]* basic_any basic_object_1))
+basic_number_7 ::= (("0") | ([1-9] [0-9]*))
+triggered_tags_group ::= (("" root_1 "end"))
+triggered_tags ::= TagDispatch(
+  ("start", triggered_tags_group),
+  stop_eos=true,
+  stop_str=(),
+  loop_after_dispatch=true
+)
+root_1_1 ::= ((triggered_tags))
+root_2 ::= (([a-z] root_2) | ([a-z]))
+"""
+    start = "start"
+    schema = {"type": "object", "properties": {"arg": {"type": "string"}}}
+    end = "end"
+    tag = xgr.StructuralTagItem(begin=start, schema=schema, end=end)
+    triggers = [start]
+    stag_grammar = xgr.Grammar.from_structural_tag([tag], triggers)
+    start_grammar = xgr.Grammar.from_ebnf("root ::= [a-z] root | [a-z]")
+    grammar_union = xgr.Grammar.union(stag_grammar, start_grammar)
+    assert str(grammar_union) == expected_grammar_union
+    grammar_concat = xgr.Grammar.concat(stag_grammar, start_grammar)
+    assert str(grammar_concat) == expected_grammar_concat
+
+
 if __name__ == "__main__":
     pytest.main(sys.argv)