Reduce state continuation hash allocations

CMakeLists.txt

@@ -80,6 +80,11 @@ list(FILTER XGRAMMAR_SOURCES_PATH EXCLUDE REGEX "${PROJECT_SOURCE_DIR}/cpp/tvm_f
 add_library(xgrammar STATIC ${XGRAMMAR_SOURCES_PATH})
 target_include_directories(xgrammar PUBLIC include)
 target_include_directories(xgrammar SYSTEM PUBLIC ${XGRAMMAR_INCLUDE_PATH})
+include(CheckCXXCompilerFlag)
+check_cxx_compiler_flag("-fno-semantic-interposition" XGRAMMAR_HAS_NO_SEMANTIC_INTERPOSITION)
+if(XGRAMMAR_HAS_NO_SEMANTIC_INTERPOSITION)
+  target_compile_options(xgrammar PRIVATE -fno-semantic-interposition)
+endif()
 
 # link to cpptrace
 if(XGRAMMAR_ENABLE_CPPTRACE)

cpp/compiled_grammar.cc

@@ -5,6 +5,8 @@
 
 #include <xgrammar/compiler.h>
 
+#include <algorithm>
+
 #include "compiled_grammar_impl.h"
 #include "support/json_serializer.h"
 #include "testing.h"
@@ -42,6 +44,28 @@ AdaptiveTokenMask::AdaptiveTokenMask(
   }
 
   this->uncertain_indices = uncertain_indices;
+  RebuildDerivedData(vocab_size, sorted_decoded_vocab);
+}
+
+AdaptiveTokenMask::AdaptiveTokenMask(
+    size_t vocab_size,
+    const std::vector<std::pair<int32_t, std::string>>& sorted_decoded_vocab,
+    std::vector<int32_t>&& accepted_indices,
+    std::vector<int32_t>&& uncertain_indices
+) {
+  const auto size_acc = accepted_indices.size();
+
+  store_type = size_acc >= USE_BITSET_THRESHOLD ? StoreType::kAcceptedBitset : StoreType::kAccepted;
+  if (store_type == StoreType::kAcceptedBitset) {
+    accepted_bitset = DynamicBitset(vocab_size);
+    for (auto idx : accepted_indices) {
+      accepted_bitset.Set(sorted_decoded_vocab[idx].first, true);
+    }
+  } else {
+    this->accepted_indices = std::move(accepted_indices);
+  }
+  this->uncertain_indices = std::move(uncertain_indices);
+  RebuildDerivedData(vocab_size, sorted_decoded_vocab);
 }
 
 AdaptiveTokenMask::AdaptiveTokenMask(
@@ -64,6 +88,39 @@ AdaptiveTokenMask::AdaptiveTokenMask(
     this->accepted_indices = accepted_indices;
   }
   this->uncertain_indices = uncertain_indices;
+  RebuildDerivedData(vocab_size, sorted_decoded_vocab);
+}
+
+void AdaptiveTokenMask::RebuildDerivedData(
+    size_t vocab_size, const std::vector<std::pair<int32_t, std::string>>& sorted_decoded_vocab
+) {
+  max_uncertain_token_length = 0;
+  if (uncertain_indices.size() < UNCERTAIN_BITSET_THRESHOLD) {
+    uncertain_bitset = DynamicBitset();
+    uncertain_lcp_with_previous.clear();
+    return;
+  }
+  uncertain_bitset = DynamicBitset(vocab_size);
+  uncertain_lcp_with_previous.resize(uncertain_indices.size());
+  const std::string* previous_token = nullptr;
+  for (size_t i = 0; i < uncertain_indices.size(); ++i) {
+    auto idx = uncertain_indices[i];
+    uncertain_bitset.Set(sorted_decoded_vocab[idx].first);
+    const auto& token = sorted_decoded_vocab[idx].second;
+    max_uncertain_token_length =
+        std::max(max_uncertain_token_length, static_cast<int32_t>(token.size()));
+    uncertain_lcp_with_previous[i] =
+        previous_token == nullptr
+            ? 0
+            : static_cast<int32_t>(
+                  std::mismatch(
+                      token.begin(), token.end(), previous_token->begin(), previous_token->end()
+                  )
+                      .first -
+                  token.begin()
+              );
+    previous_token = &token;
+  }
 }
 
 std::string AdaptiveTokenMask::Print(const TokenizerInfo& tokenizer_info) const {
@@ -168,7 +225,17 @@ picojson::value SerializeJSONValue(const CompiledGrammar::Impl& impl) {
   auto result = picojson::object{};
   result["grammar"] = AutoSerializeJSONValue(impl.grammar);
   result["tokenizer_metadata"] = impl.tokenizer_info->DumpMetadataValue();
-  result["adaptive_token_mask_cache"] = AutoSerializeJSONValue(impl.adaptive_token_mask_cache);
+  // Preserve the v14 wire format: serialized grammars store one mask value per ParserState even
+  // though the in-memory representation shares structurally identical masks.
+  std::unordered_map<ParserState, AdaptiveTokenMask, StateHashForCache, StateEqualForCache>
+      serialized_adaptive_token_mask_cache;
+  serialized_adaptive_token_mask_cache.reserve(impl.adaptive_token_mask_ids.size());
+  for (const auto& [state, mask_id] : impl.adaptive_token_mask_ids) {
+    XGRAMMAR_DCHECK(mask_id < impl.adaptive_token_masks.size());
+    serialized_adaptive_token_mask_cache.emplace(state, impl.adaptive_token_masks[mask_id]);
+  }
+  result["adaptive_token_mask_cache"] =
+      AutoSerializeJSONValue(serialized_adaptive_token_mask_cache);
   return picojson::value(result);
 }
 
@@ -186,6 +253,7 @@ std::optional<SerializationError> DeserializeJSONValue(
     return ConstructDeserializeError("Expect a 'grammar' field", type_name);
   }
   AutoDeserializeJSONValue(&(impl->grammar), object["grammar"], type_name);
+  impl->earley_parser_metadata = EarleyParserGrammarMetadata(impl->grammar);
   if (object.find("tokenizer_metadata") == object.end()) {
     return ConstructDeserializeError("Expect a 'tokenizer_metadata' field", type_name);
   }
@@ -199,14 +267,29 @@ 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"]);
+  std::unordered_map<ParserState, AdaptiveTokenMask, StateHashForCache, StateEqualForCache>
+      serialized_adaptive_token_mask_cache;
+  AutoDeserializeJSONValue(
+      &serialized_adaptive_token_mask_cache, object["adaptive_token_mask_cache"]
+  );
+  const auto& sorted_decoded_vocab = tokenizer_info.GetSortedDecodedVocab();
+  const auto vocab_size = tokenizer_info.GetVocabSize();
+  impl->adaptive_token_masks.reserve(serialized_adaptive_token_mask_cache.size());
+  impl->adaptive_token_mask_ids.reserve(serialized_adaptive_token_mask_cache.size());
+  for (auto& [state, mask] : serialized_adaptive_token_mask_cache) {
+    mask.RebuildDerivedData(vocab_size, sorted_decoded_vocab);
+    const uint32_t mask_id = static_cast<uint32_t>(impl->adaptive_token_masks.size());
+    impl->adaptive_token_masks.push_back(std::move(mask));
+    impl->adaptive_token_mask_ids.emplace(state, mask_id);
+  }
   return std::nullopt;
 }
 
 /************** CompiledGrammar **************/
 
 std::size_t MemorySize(const CompiledGrammar::Impl& impl) {
-  return MemorySize(impl.grammar) + MemorySize(impl.adaptive_token_mask_cache);
+  return MemorySize(impl.grammar) + MemorySize(impl.earley_parser_metadata) +
+         MemorySize(impl.adaptive_token_masks) + MemorySize(impl.adaptive_token_mask_ids);
 }
 
 std::size_t CompiledGrammar::MemorySizeBytes() const { return MemorySize(*pimpl_); }

cpp/compiled_grammar_impl.h

@@ -52,12 +52,17 @@ struct AdaptiveTokenMask {
   StoreType store_type;
 
   static constexpr int USE_BITSET_THRESHOLD = 1000;
+  static constexpr int UNCERTAIN_BITSET_THRESHOLD = 1024;
 
   std::vector<int32_t> accepted_indices;
   std::vector<int32_t> rejected_indices;
   DynamicBitset accepted_bitset;
 
   std::vector<int32_t> uncertain_indices;
+  // Derived, non-serialized data for batching large runtime uncertain-token operations.
+  DynamicBitset uncertain_bitset;
+  std::vector<int32_t> uncertain_lcp_with_previous;
+  int32_t max_uncertain_token_length = 0;
 
   /*! \brief Default constructor. Only for deserialization. */
   AdaptiveTokenMask() = default;
@@ -77,11 +82,23 @@ struct AdaptiveTokenMask {
       const std::vector<int32_t>& uncertain_indices
   );
 
+  AdaptiveTokenMask(
+      size_t vocab_size,
+      const std::vector<std::pair<int32_t, std::string>>& sorted_decoded_vocab,
+      std::vector<int32_t>&& accepted_indices,
+      std::vector<int32_t>&& uncertain_indices
+  );
+
+  void RebuildDerivedData(
+      size_t vocab_size, const std::vector<std::pair<int32_t, std::string>>& sorted_decoded_vocab
+  );
+
   std::string Print(const TokenizerInfo& tokenizer_info) const;
 
   friend std::size_t MemorySize(const AdaptiveTokenMask& mask) {
     return MemorySize(mask.uncertain_indices) + MemorySize(mask.accepted_indices) +
-           MemorySize(mask.rejected_indices) + MemorySize(mask.accepted_bitset);
+           MemorySize(mask.rejected_indices) + MemorySize(mask.accepted_bitset) +
+           MemorySize(mask.uncertain_bitset) + MemorySize(mask.uncertain_lcp_with_previous);
   }
 };
 
@@ -99,6 +116,32 @@ XGRAMMAR_MEMBER_TABLE(
     &AdaptiveTokenMask::uncertain_indices
 );
 
+#ifdef XGRAMMAR_PROFILE_COMPILE
+/*!
+ * \brief Profile-only cost of computing one structurally deduplicated adaptive-mask task.
+ *
+ * Several ParserStates can share one task result. Runtime coverage is therefore measured both by
+ * state (stored mask bytes) and by task group (compile work that a lazy compiler would execute).
+ */
+struct AdaptiveMaskCompileProfile {
+  ParserState representative_state;
+  uint64_t mask_cpu_us = 0;
+  uint64_t possible_tokens = 0;
+  uint64_t rule_cache_hits = 0;
+  uint64_t first_byte_cache_reused_tokens = 0;
+  uint64_t saturated_cache_reused_tokens = 0;
+  uint64_t token_edge_skipped_tokens = 0;
+  uint64_t speculative_accepted_tokens = 0;
+  uint64_t subtree_pruned_tokens = 0;
+  uint64_t parser_simulated_tokens = 0;
+  uint64_t parser_naive_token_bytes = 0;
+  uint64_t parser_advance_calls = 0;
+  uint64_t parser_failed_advance_calls = 0;
+  uint64_t mask_bytes = 0;
+  uint64_t state_count = 0;
+};
+#endif
+
 /*!
  * \brief All information that we need to match tokens in the tokenizer to the specified grammar.
  * It is the result of preprocessing.
@@ -112,12 +155,32 @@ class CompiledGrammar::Impl {
   /*! \brief The tokenizer information. */
   TokenizerInfo tokenizer_info{NullObj{}};
 
+  /*! \brief Grammar-only Earley metadata shared by compile-time tasks and runtime matchers. */
+  EarleyParserGrammarMetadata earley_parser_metadata;
+
   /*! \brief Default constructor. */
   Impl() = default;
 
-  /*! \brief Mapping from the parser state to the adaptive token mask. */
-  std::unordered_map<ParserState, AdaptiveTokenMask, StateHashForCache, StateEqualForCache>
-      adaptive_token_mask_cache;
+  /*! \brief Structurally deduplicated adaptive token masks. */
+  std::vector<AdaptiveTokenMask> adaptive_token_masks;
+
+  /*! \brief Mapping from each parser state to an entry in adaptive_token_masks. */
+  std::unordered_map<ParserState, uint32_t, StateHashForCache, StateEqualForCache>
+      adaptive_token_mask_ids;
+
+#ifdef XGRAMMAR_PROFILE_COMPILE
+  /*! \brief Per-task compile work, indexed by the structurally deduplicated task-group id. */
+  std::vector<AdaptiveMaskCompileProfile> adaptive_mask_compile_profiles;
+#endif
+
+  const AdaptiveTokenMask* FindAdaptiveTokenMask(const ParserState& state) const {
+    const auto mask_id_it = adaptive_token_mask_ids.find(state);
+    if (mask_id_it == adaptive_token_mask_ids.end()) {
+      return nullptr;
+    }
+    XGRAMMAR_DCHECK(mask_id_it->second < adaptive_token_masks.size());
+    return &adaptive_token_masks[mask_id_it->second];
+  }
 
   Grammar GetGrammar() const { return grammar; }
 
@@ -139,8 +202,10 @@ XGRAMMAR_MEMBER_TABLE(
     &CompiledGrammar::Impl::grammar,
     "tokenizer_info",
     &CompiledGrammar::Impl::tokenizer_info,
-    "adaptive_token_mask_cache",
-    &CompiledGrammar::Impl::adaptive_token_mask_cache
+    "adaptive_token_masks",
+    &CompiledGrammar::Impl::adaptive_token_masks,
+    "adaptive_token_mask_ids",
+    &CompiledGrammar::Impl::adaptive_token_mask_ids
 );
 
 }  // namespace xgrammar