[Spec][Ngram] 5/N: Store and advance anchor match state across decode steps (#21243)

python/sglang/jit_kernel/ngram_corpus.py

@@ -85,4 +85,31 @@ def get_ngram_corpus_cls():
                 np.int64
             )
 
+        def match_stateful(
+            self,
+            state_ids: List[int],
+            batch_tokens: List[List[int]],
+            total_lens: List[int],
+        ) -> Tuple[np.ndarray, np.ndarray]:
+            tokens_flat, offsets = _to_csr(batch_tokens)
+            batch_size = len(batch_tokens)
+            d = self._draft_token_num
+
+            state_ids_t = torch.tensor(state_ids, dtype=torch.int64)
+            total_lens_t = torch.tensor(total_lens, dtype=torch.int64)
+            out_tokens = torch.zeros(batch_size * d, dtype=torch.int32)
+            out_mask = torch.zeros(batch_size * d * d, dtype=torch.uint8)
+
+            self.batch_match_stateful(  # type: ignore
+                state_ids_t, tokens_flat, offsets, total_lens_t, out_tokens, out_mask
+            )
+
+            return out_tokens.numpy().astype(np.int64), out_mask.numpy().astype(
+                np.int64
+            )
+
+        def erase_states(self, state_ids: List[int]) -> None:
+            state_ids_t = torch.tensor(state_ids, dtype=torch.int64)
+            self.erase_match_state(state_ids_t)  # type: ignore
+
     return NgramCorpusFFI