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

test/registered/unit/spec/test_ngram_corpus.py

@@ -4,10 +4,10 @@ import uuid
 import numpy as np
 
 from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=8, suite="stage-b-test-1-gpu-small")
+register_cpu_ci(est_time=10, suite="stage-a-test-cpu")
 
 
 def _make_corpus(match_type="BFS", **kwargs):
@@ -674,5 +674,69 @@ class TestNgramCorpusIncremental(CustomTestCase):
         np.testing.assert_array_equal(inc_masks, full_masks)
 
 
+class TestNgramCorpusMatchBenchmark(CustomTestCase):
+    """Benchmark incremental advance vs full rebuild in match()."""
+
+    def test_incremental_faster_than_rebuild(self):
+        """Incremental advance (O(D) per token) should be faster than rebuild (O(D^2))."""
+        import time
+
+        max_trie_depth = 18
+        draft_token_num = 8
+        corpus = _make_corpus(
+            "BFS",
+            max_trie_depth=max_trie_depth,
+            draft_token_num=draft_token_num,
+            capacity=500000,
+        )
+
+        # Seed the trie with diverse sequences so suffix matching is non-trivial.
+        seed_data = [list(range(i, i + 50)) for i in range(0, 5000, 50)]
+        corpus.batch_put(seed_data)
+        corpus.synchronize()
+
+        num_steps = 500
+        base_seq = list(range(1, max_trie_depth + 1))
+
+        # --- Incremental path: same req_id, total_len grows by 1 each step ---
+        req_id = "bench-incremental"
+        # Warm up the state with the initial context.
+        _batch_get_with_state(corpus, req_id, base_seq, len(base_seq))
+
+        start = time.perf_counter()
+        for step in range(num_steps):
+            total_len = len(base_seq) + step + 1
+            new_token = (step + max_trie_depth + 1) % 5000
+            tail = (base_seq + [new_token])[-max_trie_depth:]
+            base_seq = tail
+            _batch_get_with_state(corpus, req_id, tail, total_len)
+        incremental_us = (time.perf_counter() - start) / num_steps * 1e6
+
+        # --- Rebuild path: unique req_id each call forces fresh state ---
+        base_seq = list(range(1, max_trie_depth + 1))
+        start = time.perf_counter()
+        for step in range(num_steps):
+            new_token = (step + max_trie_depth + 1) % 5000
+            tail = (base_seq + [new_token])[-max_trie_depth:]
+            base_seq = tail
+            _batch_get(corpus, [tail])
+        rebuild_us = (time.perf_counter() - start) / num_steps * 1e6
+
+        print(
+            f"\n  Incremental: {incremental_us:.1f} us/step"
+            f"\n  Rebuild:     {rebuild_us:.1f} us/step"
+            f"\n  Speedup:     {rebuild_us / incremental_us:.2f}x"
+        )
+
+        # The incremental path should be at least as fast; allow a small margin
+        # for noise. With D=12 the theoretical speedup is ~12x (D^2/D).
+        self.assertLess(
+            incremental_us,
+            rebuild_us * 1.1,
+            f"Incremental ({incremental_us:.1f} us) should not be slower than "
+            f"rebuild ({rebuild_us:.1f} us)",
+        )
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=3)