[EtsyML] Native TRIE guided decoding (per-beam, host-staging)#1
Open
orsonadams wants to merge 8 commits into
Open
[EtsyML] Native TRIE guided decoding (per-beam, host-staging)#1orsonadams wants to merge 8 commits into
orsonadams wants to merge 8 commits into
Conversation
- Add GuidedTrie and GuidedTrieNode support in guidedDecoder for trie-based token constraint handling with bitmask caching - Extend GuidedDecodingConfig with trie_path parameter in executor configs - Add Python bindings for trie_path in nanobind and pybind executor configs - Add trie_decoding module with tools for trie generation, conversion, load testing, and documentation - Add trie_loader, trie_mask, and trie_mask_xgr libraries for Triton backend - Update Triton model.py to integrate trie-based decoding masks - Update Triton config.pbtxt files with trie_path parameters
Replace the per-node unordered_map<int32,shared_ptr<GuidedTrieNode>> trie with a flat CSR representation (child offsets + token-sorted child arrays + isEnd), and allow the trie EOS token at terminal nodes so explicit '...,EOS' leaf nodes are no longer materialized. Together these cut runtime memory for the gen_search 73M-SID trie from ~20GB to ~1GB (no per-node heap/hashmap; ~half the node count). Binary format gains an eos_token_id header field; convert_trie_to_binary.py writes it and no longer needs EOS appended to sequences. Also demote per-step TRIE debug logging from INFO to DEBUG. Verified: gen_search BART enc-dec produces the identical constrained SID as the pointer-based trie. Signed-off-by: Orson Adams <oadams@etsy.com>
Teach the tensorrtllm C++ Triton backend (libtriton_tensorrtllm.so) to construct a GuidedDecodingConfig(kTRIE) from config.pbtxt when guided_decoding_backend=trie, reading the binary trie path from the trie_path parameter. Previously TRIE was only wired in the Python-backend model.py; prod serves via the C++ backend for throughput (beam search), so the C++ path needs it. Unlike xgrammar, TRIE needs only the trie path (no tokenizer info / encoded vocab). Signed-off-by: Orson Adams <oadams@etsy.com>
Replace the per-node v1 parse loop with a v2 format that stores the CSR arrays contiguously (magic, version, numNodes, numEdges, eos, childOffsets, childTokens, childNodes, isEnd). Loading becomes four bulk reads (O(file size), ~6s for the 73M-SID gen_search trie) and the loader validates the format version and truncation. The builder writes children token-sorted, so the runtime-side child sort is removed. Python writer updated to v2; v1 files are rejected with a clear error. Signed-off-by: Orson Adams <oadams@etsy.com>
…eam search) The TRIE backend previously kept one trie state per request (advanced by beam 0's last token) and masked only beam 0's logits row, so at beam>1 only the first token was constrained and outputs were invalid. Also, request token buffers are raw slot buffers whose history goes stale when beam search switches hypothesis parents, so any re-walk from them is unsound. Redesign: - GuidedDecoder::build re-walks every beam's generated tokens from the trie root each step (stateless, bounded by trie depth) and stages one bitmask row per (seqSlot, beam); off-trie or finished beams get an EOS-only mask so they terminate instead of sampling all -inf logits. - GuidedDecoder::execute registers beamWidth (logits row, bitmask row) entries per request; the logitsBitmask kernel is unchanged. - Before build(), the model refreshes each beam-search request's token histories with the same non-destructive per-step gatherTree used by streaming+beam (postProcessRequest gains a streamingGather flag), skipped until the first token exists. - Bitmask buffers grow to [maxNumSequences * maxBeamWidth, bitmaskSize]; the xgrammar backend keeps its per-request design on beam 0's row. Validated on the gen_search BART enc-dec engine with the full 73M-SID trie: beam=128 yields 128/128 valid unique SIDs per query (was 0/128); beam=1 unchanged. Per-query latency at beam=128: 27.9ms vs 12.8ms unconstrained. Signed-off-by: Orson Adams <oadams@etsy.com>
…ild context cpp/build, .git, *.whl and trie_decoding data are not needed by the image build and inflate the build context by tens of GB. Signed-off-by: Orson Adams <oadams@etsy.com>
The per-beam TRIE mask path issued two hard stream synchronizations every decode step: one after the pointer-vector H2D copies and one after the logits-bitmask kernel. The copies, the kernel, and the downstream decoder all run on the runtime stream, so stream ordering already guarantees the masks land before the logits are consumed -- the syncs only stalled the host. Remove both (plus the now-misleading "kernel completed" debug lines). Correctness unchanged (beam=128 640/640 valid/128 unique, beam=1 5/5, beam_threshold 17/128 OK). Trie overhead at beam=128 drops ~1.7 ms/query (27.9 -> 26.2), ~0.45 ms/decode-step. Signed-off-by: Orson Adams <oadams@etsy.com>
Rename the native-TRIE config.pbtxt parameter the C++ backend reads from "trie_path" to "guided_decoding_trie_path", pairing consistently with the existing "guided_decoding_backend" param. The etsyml model manifest and the kserve handler emit this key; keeping the name aligned end to end. Signed-off-by: Orson Adams <oadams@etsy.com>
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.
Native trie-based constrained beam search in the C++
GuidedDecoder, for gen_search on the Tritontensorrt_llmbackend. Base is clean upstream 1.3.0rc21 (d567f924b7) so this diff is only the fork's trie work.Contents
GuidedDecoderkTRIE path.guided_decoding_backend+guided_decoding_trie_pathconfig params.execute().Beam-correct at beam=128 (640/640 valid, 128 unique). The device-resident kernel that removes the per-step host work stacks on top of this in a follow-up PR.
Why a fork (and why not the Python
LogitsPostProcessor): https://docs.google.com/document/d/1OsciNfXKMkCgfpI-8jkN7qo_yeLnidknRbGQPvCIBlw/edit