perf: speed up guided decoding with xgrammar new version and batched update

Author: windreamer

CMakeLists.txt

@@ -100,7 +100,7 @@ FetchContent_MakeAvailable(yaml-cpp)
 FetchContent_Declare(
   xgrammar
   GIT_REPOSITORY          https://github.com/mlc-ai/xgrammar.git
-  GIT_TAG                 v0.1.27
+  GIT_TAG                 v0.2.1
   GIT_SUBMODULES          "3rdparty/dlpack"
   GIT_PROGRESS            TRUE
   USES_TERMINAL_DOWNLOAD  TRUE

lmdeploy/turbomind/turbomind.py

@@ -165,6 +165,7 @@ def __init__(self,
                                                       trust_remote_code=trust_remote_code)
         self.is_dummy = self.model_comm.is_dummy_node()
         self.tokenizer = Tokenizer(model_path, trust_remote_code=trust_remote_code)
+        self._grammar_compiler = None
         if not _engine_config.empty_init:
             with torch.cuda.device(self.devices[0]):
                 model_loader.export()
@@ -173,6 +174,14 @@ def __init__(self,
 
         self.session_len = _engine_config.session_len
 
+    @property
+    def grammar_compiler(self):
+        """Lazy-initialized GrammarCompiler shared across all requests."""
+        if self._grammar_compiler is None:
+            tokenizer_info = TokenizerInfo.from_huggingface(
+                self.tokenizer.model.model, vocab_size=self._vocab_size)
+            self._grammar_compiler = _xgr.GrammarCompiler(tokenizer_info)
+        return self._grammar_compiler
 
     def _process_weights(self):
         """Process weight."""
@@ -672,11 +681,8 @@ async def async_stream_infer(self,
                                                 gen_config=gen_config)
 
         if gen_config.response_format is not None:
-            tokenizer = self.tm_model.tokenizer
-            vocab_size = self.tm_model._vocab_size
-
             try:
-                tokenizer_info = TokenizerInfo.from_huggingface(tokenizer.model.model, vocab_size=vocab_size)
+                compiler = self.tm_model.grammar_compiler
                 decode_grammar_type = gen_config.response_format['type']
                 if decode_grammar_type == 'json_schema':
                     decode_grammar = gen_config.response_format[decode_grammar_type]['schema']
@@ -685,8 +691,6 @@ async def async_stream_infer(self,
                 elif decode_grammar_type == 'json_object':
                     decode_grammar = '{"type" : "object", "additionalProperties": true}'
 
-                compiler = _xgr.GrammarCompiler(tokenizer_info)
-
                 if decode_grammar_type == 'json_schema':
                     decode_grammar = json.dumps(decode_grammar)
                     grammar = compiler.compile_json_schema(decode_grammar)
@@ -702,7 +706,7 @@ async def async_stream_infer(self,
 
                 self.model_inst.set_grammar(grammar)
             except ValueError as e:
-                logger.warning(f'Failed to initialize guided decoding for tokenizer {tokenizer}, '
+                logger.warning(f'Failed to initialize guided decoding, '
                                f'disable guided decoding: {e}')
                 gen_config.response_format = None
 

src/turbomind/generation/guided_decoding.cc

@@ -12,6 +12,7 @@ namespace turbomind {
 struct GuidedDecoding::Data {
     Tensor_<int32_t> bitmask;
     bool             active{};
+    bool             needs_apply{};
 
     std::vector<std::shared_ptr<xgrammar::GrammarMatcher>> matchers;
 };
@@ -56,24 +57,36 @@ void GuidedDecoding::FillMask(int phase, TensorMap& env)
                            (int64_t*)bitmask_buf_.shape().data(),
                            nullptr,
                            0};
+
+        std::vector<xgrammar::GrammarMatcher> active_matchers;
+        std::vector<int32_t>                  active_indices;
+
+        d.needs_apply = false;
+
         if (tp_group_->rank() == 0) {
             for (size_t i = 0; i < d.matchers.size(); ++i) {
-                if (const auto& matcher = d.matchers[i]; matcher && !matcher->IsTerminated()) {
-                    matcher->FillNextTokenBitmask(&dlbitmask, i);
+                if (const auto& m = d.matchers[i]; m && !m->IsTerminated()) {
+                    active_matchers.push_back(*m);
+                    active_indices.push_back(static_cast<int32_t>(i));
                 }
                 else {
                     std::fill_n(bitmask_buf_.data() + i * bitmask_buf_.stride(0),
                                 bitmask_buf_.stride(0),
                                 static_cast<int32_t>(-1));
                 }
             }
+
+            if (!active_matchers.empty()) {
+                batch_matcher_.BatchFillNextTokenBitmask(&active_matchers, &dlbitmask, active_indices);
+                d.needs_apply = true;
+            }
         }
     }
 }
 
 void GuidedDecoding::ApplyMask(int phase, TensorMap& env)
 {
-    if (auto& d = *data_.at(phase); d.active) {
+    if (auto& d = *data_.at(phase); d.active && d.needs_apply) {
         const ssize_t numel = d.matchers.size() * bitmask_buf_.stride(0);
         if (tp_group_->n_ranks() > 1) {
             // bcast the data instead of `bitmask_buf` instance (which may avoid copying the data)
@@ -93,13 +106,24 @@ void GuidedDecoding::Update(int phase, TensorMap& env)
     if (auto& d = *data_.at(phase); d.active) {
         Copy(env.at("output_ids").buffer(), d.matchers.size(), output_ids_buf_);
         core::Context::stream().Sync();
+
         if (tp_group_->rank() == 0) {
+            // Collect active matchers and their token IDs for batch AcceptToken
+            std::vector<xgrammar::GrammarMatcher> active_matchers;
+            std::vector<int32_t>                  active_token_ids;
+
             for (size_t i = 0; i < d.matchers.size(); ++i) {
-                if (const auto& matcher = d.matchers[i]; matcher && !matcher->IsTerminated()) {
-                    matcher->AcceptToken(output_ids_buf_[i]);
+                if (const auto& m = d.matchers[i]; m && !m->IsTerminated()) {
+                    active_matchers.push_back(*m);  // copy: refcount++
+                    active_token_ids.push_back(output_ids_buf_[i]);
                 }
             }
+
+            if (!active_matchers.empty()) {
+                xgrammar::BatchGrammarMatcher::BatchAcceptToken(&active_matchers, active_token_ids);
+            }
         }
+        // active_matchers destroyed here: refcount-- for each entry
     }
 }
 

src/turbomind/generation/guided_decoding.h

@@ -6,6 +6,7 @@
 
 #include "src/turbomind/comm/host_comm.h"
 #include "src/turbomind/core/core.h"
+#include "xgrammar/matcher.h"
 
 namespace turbomind {
 
@@ -27,6 +28,8 @@ class GuidedDecoding: public BaseGenerationParam {
     struct Data;
     std::vector<std::shared_ptr<Data>> data_;
 
+    xgrammar::BatchGrammarMatcher batch_matcher_;
+
     Tensor_<int32_t> bitmask_buf_;
     Buffer_<int>     output_ids_buf_;
 };