[Spec][Ngram] Followup fixes for MatchState incremental advance (#22180)

python/sglang/jit_kernel/csrc/ngram_corpus/trie.cpp

@@ -91,6 +91,8 @@ void Trie::squeeze(size_t count) {
 }
 
 void Trie::reset() {
+  // Epoch bump invalidates all cached MatchState objects, so we do not need to
+  // retireNode() on every node individually.
   ++trie_epoch_;
   global_lru_.clear();
   path_.clear();
@@ -150,16 +152,17 @@ bool Trie::advanceMatchState_(MatchState& state, const int32_t* tokens, size_t l
     return false;
   }
 
+  // Reuse a single buffer across iterations to avoid per-token heap allocation.
+  std::vector<NodeRef> next;
+  next.reserve(param_.max_trie_depth);
+
   for (size_t i = 0; i < len; ++i) {
     const auto next_depth = std::min(state.anchors.size() + 1, param_.max_trie_depth);
-    std::vector<NodeRef> next(next_depth);
+    next.assign(next_depth, {});
 
-    const auto root_ref = rootRef();
-    const auto root = resolve(state, root_ref);
-    if (root == nullptr) {
-      return false;
-    }
-    if (auto iter = root->child.find(tokens[i]); iter != root->child.end()) {
+    // Root is never evicted, so we access it directly; the epoch was already
+    // validated above.
+    if (auto iter = root_->child.find(tokens[i]); iter != root_->child.end()) {
       next[0] = capture(iter->second);
     }
 

python/sglang/jit_kernel/csrc/ngram_corpus/trie.h

@@ -24,6 +24,8 @@ struct TrieNode {
   int32_t freq = 0;
   // Logical generation of this TrieNode. retireNode() bumps it before the node
   // goes back to the pool so stale NodeRefs fail validation after reuse.
+  // Starts at 1 so that a default-constructed NodeRef (version=0) never
+  // accidentally resolves to a live node.
   uint64_t version = 1;
 
   struct CompareByFreq {

python/sglang/srt/speculative/ngram_worker.py

@@ -269,6 +269,10 @@ class NGRAMWorker:
             if batch.return_logprob:
                 add_output_logprobs_for_spec_v1(batch, verify_input, logits_output)
             self._update_ngram_corpus(batch)
+            # Clean up per-request match state for finished/retracted requests.
+            # State entries are created in _prepare_draft_tokens and cleaned here.
+            # If a request is removed without passing through verify, the entry
+            # persists until reset(); this is acceptable because MatchState is small.
             finished_req_ids = []
             for req in batch.reqs:
                 if req.finished() or req.is_retracted: