[Spec][Ngram] 5/N: Store and advance anchor match state across decode steps#21243
Merged
hnyls2002 merged 8 commits intosgl-project:mainfrom Apr 6, 2026
Conversation
kpham-sgl
requested review from
Ying1123,
hnyls2002 and
merrymercy
as code owners
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
github-actions
bot
added
documentation
This was referenced Mar 24, 2026
…-decode-steps Resolve merge conflicts adapting the stateful ngram MatchState feature to the upstream's TVM FFI binding (replacing pybind11): - trie.h/trie.cpp: keep HEAD's MatchState, NodeRef, and incremental anchor advancement logic (auto-resolved) - ngram.h/ngram.cpp: add stateless batchMatch(tokens) overload for FFI, switch stateful overload to int64_t keys (FFI-compatible) - ngram_corpus_ffi.cpp: add batch_match_stateful and erase_match_state FFI methods - jit_kernel/ngram_corpus.py: add match_stateful and erase_states to the FFI wrapper - srt/speculative/cpp_ngram/ngram_corpus.py: use get_ngram_corpus_cls() with a req_id-to-state_id mapping layer for the stateful path - Remove old pybind11 ngram_corpus_binding.cpp (superseded by FFI) - test_ngram_corpus.py: resolve to use _batch_get helper, remove stale _raw_batch_match Made-with: Cursor
kpham-sgl
requested review from
BBuf,
DarkSharpness,
HydraQYH,
celve and
yuan-luo
as code owners
Collaborator
|
/rerun-test test_ngram_corpus test_ngram_speculative_decoding |
Contributor
|
❌ ❌ |
Collaborator
|
/rerun-test test/registered/spec/utils/test_ngram_corpus.py test/registered/spec/test_ngram_speculative_decoding.py |
Contributor
|
✅ ✅ |
Collaborator
|
/tag-and-rerun-ci |
hnyls2002
added a commit
to kpham-sgl/sglang
that referenced
this pull request
Apr 6, 2026
Adapt SAM external corpus feature to the new ngram architecture:
- Move suffix_automaton.{h,cpp} to jit_kernel/csrc/ngram_corpus/
- Use TVM FFI instead of pybind11 for SAM loading methods
- Change match_state_ key from std::string to int64_t state_ids
- Add SAM budget splitting to stateful batchMatch overload
- Wire external corpus params through FFI constructor
- Keep both stateless and stateful batchMatch overloads
JustinTong0323
pushed a commit
to JustinTong0323/sglang
that referenced
this pull request
Apr 7, 2026
Fridge003
pushed a commit
that referenced
this pull request
Apr 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Part of Ngram refactoring series #21052
Following #21225
Previously, match() is O(D^2) at every decode steps where D is the
max_trie_depth. One observation is, since we always append to the sequence during decode, we can store a list ofMatchState(which corresponds to an anchor) from previous decode step and advance them in O(1) for each anchorModifications
MatchStateplus versionedNodeRefso cached anchors can be advanced safely across decode steps and invalidated correctly after eviction or reset.Trie::match()stateful: infer the appended suffix from the current tail andtotal_len, advance cached anchors when valid, and rebuild when state is stale.Ngram, keyed byreq_id, with explicit cleanup on request finish/reset.batchMatch(req_ids, tokens, total_lens)/batch_get(req_ids, batch_tokens, total_lens)by removing explicit appended-token plumbing.NGRAMWorkerintegration so it only passes the trimmed tail and full request length, without mutatingReqor keeping extra Python-side match state.Accuracy Tests
Passed
python3 -m pytest -q test/registered/spec/utils/test_ngram_corpus.pyandpython3 -m pytest -q test/registered/spec/test_ngram_speculative_decoding.pyBenchmarking and Profiling
[TODO]
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci