Add TraverseDraftTree for speculative decoding bitmask generation - #490
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a new TraverseDraftTree function that performs depth-first traversal of a speculative decoding tree to generate grammar-aware token bitmasks for each node in a single pass. This enables efficient validation of multiple speculative tokens against grammar constraints.
Key changes:
- Implements core C++ DFS traversal algorithm with grammar validation and rollback logic
- Adds Python bindings and wrapper function for the new tree traversal functionality
- Includes minor formatting improvements to test files
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| cpp/testing.h | Adds function declaration and documentation for TraverseDraftTree |
| cpp/testing.cc | Implements TraverseDraftTree and internal DFS helper for tree traversal with grammar matching |
| cpp/nanobind/nanobind.cc | Adds Python binding for _traverse_draft_tree with DLTensor conversion |
| python/xgrammar/testing.py | Adds Python wrapper function _traverse_draft_tree with parameter documentation |
| tests/python/test_grammar_parser_macro.py | Adds blank line for formatting consistency |
| tests/python/test_grammar_matcher_json_schema.py | Splits long JSON schema string across multiple lines for readability |
| docs/conf.py | Adds comment placeholder for documentation field overrides |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // Convert ndarrays to DLTensors | ||
| static_assert( | ||
| sizeof(retrieve_next_token) == sizeof(void*) + sizeof(nb::dlpack::dltensor) | ||
| ); |
There was a problem hiding this comment.
The static_assert here checks only the retrieve_next_token array size, but should ideally verify all the ndarray parameters have the expected layout since they all undergo the same pointer reinterpretation. Consider adding checks for retrieve_next_sibling, draft_tokens, and bitmask as well, similar to how the function at line 76 and 101 handle this pattern.
| ); | |
| ); | |
| static_assert( | |
| sizeof(retrieve_next_sibling) == sizeof(void*) + sizeof(nb::dlpack::dltensor) | |
| ); | |
| static_assert( | |
| sizeof(draft_tokens) == sizeof(void*) + sizeof(nb::dlpack::dltensor) | |
| ); | |
| static_assert( | |
| sizeof(bitmask) == sizeof(void*) + sizeof(nb::dlpack::dltensor) | |
| ); |
| .def( | ||
| "_traverse_draft_tree", | ||
| [](nb::ndarray<> retrieve_next_token, | ||
| nb::ndarray<> retrieve_next_sibling, | ||
| nb::ndarray<> draft_tokens, | ||
| GrammarMatcher& matcher, | ||
| nb::ndarray<> bitmask) { | ||
| // Convert ndarrays to DLTensors | ||
| static_assert( | ||
| sizeof(retrieve_next_token) == sizeof(void*) + sizeof(nb::dlpack::dltensor) | ||
| ); | ||
| DLTensor* next_token_ptr = reinterpret_cast<DLTensor*>( | ||
| reinterpret_cast<char*>(&retrieve_next_token) + sizeof(void*) | ||
| ); | ||
| DLTensor* next_sibling_ptr = reinterpret_cast<DLTensor*>( | ||
| reinterpret_cast<char*>(&retrieve_next_sibling) + sizeof(void*) | ||
| ); | ||
| DLTensor* draft_tokens_ptr = | ||
| reinterpret_cast<DLTensor*>(reinterpret_cast<char*>(&draft_tokens) + sizeof(void*)); | ||
| DLTensor* bitmask_ptr = | ||
| reinterpret_cast<DLTensor*>(reinterpret_cast<char*>(&bitmask) + sizeof(void*)); | ||
| TraverseDraftTree( | ||
| next_token_ptr, next_sibling_ptr, draft_tokens_ptr, matcher, bitmask_ptr | ||
| ); | ||
| }, |
There was a problem hiding this comment.
The _traverse_draft_tree function is missing input validation that other similar functions in this file have. It should validate that: 1) all input arrays are 1D (retrieve_next_token, retrieve_next_sibling, draft_tokens) or 2D (bitmask), 2) all arrays are on CPU, and 3) all arrays have int32 dtype. See the validation patterns in GrammarMatcher_FillNextTokenBitmask (lines 54-66) and GrammarMatcher_BatchFillNextTokenMask (lines 92-100) for reference.
| retrieve_next_sibling: torch.Tensor, | ||
| draft_tokens: torch.Tensor, | ||
| matcher: "GrammarMatcher", | ||
| allocate_token_bitmask: torch.Tensor, |
There was a problem hiding this comment.
The parameter name "allocate_token_bitmask" is inconsistent with the documentation and other usages. In the documentation (line 375) and throughout the codebase, this is referred to as "bitmask". The name "allocate_token_bitmask" sounds like a function name rather than a parameter for a pre-allocated bitmask. Consider renaming to just "bitmask" or "token_bitmask" for consistency.
| def _traverse_draft_tree( | ||
| retrieve_next_token: torch.Tensor, | ||
| retrieve_next_sibling: torch.Tensor, | ||
| draft_tokens: torch.Tensor, | ||
| matcher: "GrammarMatcher", | ||
| allocate_token_bitmask: torch.Tensor, | ||
| ) -> None: | ||
| """Traverse the tree constructed by the draft model to generate the logits mask. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| retrieve_next_token : torch.Tensor | ||
| 1D int32 tensor where retrieve_next_token[i] gives the index of the child node | ||
| of node i, or -1 if no child exists. | ||
| retrieve_next_sibling : torch.Tensor | ||
| 1D int32 tensor where retrieve_next_sibling[i] gives the index of the sibling node | ||
| of node i, or -1 if no sibling exists. | ||
| draft_tokens : torch.Tensor | ||
| 1D int32 tensor of draft token ids at each position in the tree. | ||
| matcher : GrammarMatcher | ||
| The grammar matcher to use for validation. | ||
| allocate_token_bitmask : torch.Tensor | ||
| 2D int32 tensor (num_nodes x bitmask_size) to store the generated bitmasks. | ||
| """ | ||
| assert retrieve_next_token.shape == retrieve_next_sibling.shape == draft_tokens.shape | ||
|
|
||
| _core.testing._traverse_draft_tree( | ||
| retrieve_next_token, | ||
| retrieve_next_sibling, | ||
| draft_tokens, | ||
| matcher._handle, | ||
| allocate_token_bitmask, | ||
| ) |
There was a problem hiding this comment.
The new _traverse_draft_tree function lacks test coverage. Given that the repository has comprehensive test coverage for other testing utility functions, this new functionality should also include tests to validate the tree traversal logic, bitmask generation, and edge cases (e.g., empty trees, single nodes, rejected tokens).
| void TraverseDraftTree( | ||
| const DLTensor* retrieve_next_token, | ||
| const DLTensor* retrieve_next_sibling, | ||
| const DLTensor* draft_tokens, | ||
| GrammarMatcher& matcher, | ||
| DLTensor* bitmask | ||
| ) { | ||
| DFS(0, | ||
| -1, | ||
| reinterpret_cast<const int32_t*>(retrieve_next_token->data), | ||
| reinterpret_cast<const int32_t*>(retrieve_next_sibling->data), | ||
| reinterpret_cast<const int32_t*>(draft_tokens->data), | ||
| matcher, | ||
| bitmask); | ||
| } |
There was a problem hiding this comment.
The TraverseDraftTree function lacks input validation. It should validate that: 1) all input DLTensors are non-null, 2) retrieve_next_token, retrieve_next_sibling, and draft_tokens have matching shapes, 3) bitmask has the expected 2D shape, and 4) all tensors have appropriate dtypes (int32). Without these checks, invalid inputs could lead to undefined behavior or crashes.
| if (curr == 0) { | ||
| // The first token generated by the target model, always accepted | ||
| accepted = true; | ||
| } else { | ||
| int32_t curr_token_id = draft_tokens[curr]; | ||
| int32_t* parent_bitmask = bitmask_data + parent_pos * bitmask_size; | ||
| // 32 boolean bitmask values are packed into 32-bit integers | ||
| accepted = (parent_bitmask[curr_token_id / 32] & (1 << (curr_token_id % 32))) != 0; |
There was a problem hiding this comment.
There's a potential out-of-bounds access vulnerability. The code reads draft_tokens[curr] without validating that curr is within the valid range of the draft_tokens tensor. Similarly, line 82 accesses parent_bitmask[curr_token_id / 32] without verifying that curr_token_id / 32 is within bounds of the bitmask_size. These accesses should be bounds-checked to prevent buffer overruns.
| if (curr == 0) { | |
| // The first token generated by the target model, always accepted | |
| accepted = true; | |
| } else { | |
| int32_t curr_token_id = draft_tokens[curr]; | |
| int32_t* parent_bitmask = bitmask_data + parent_pos * bitmask_size; | |
| // 32 boolean bitmask values are packed into 32-bit integers | |
| accepted = (parent_bitmask[curr_token_id / 32] & (1 << (curr_token_id % 32))) != 0; | |
| // Defensive: get the number of draft tokens | |
| int32_t num_draft_tokens = static_cast<int32_t>(bitmask->shape[0]); | |
| if (curr < 0 || curr >= num_draft_tokens) { | |
| // Out of bounds, skip this branch | |
| return; | |
| } | |
| if (curr == 0) { | |
| // The first token generated by the target model, always accepted | |
| accepted = true; | |
| } else { | |
| int32_t curr_token_id = draft_tokens[curr]; | |
| if (curr_token_id < 0) { | |
| // Invalid token id, skip this branch | |
| return; | |
| } | |
| int32_t* parent_bitmask = bitmask_data + parent_pos * bitmask_size; | |
| // 32 boolean bitmask values are packed into 32-bit integers | |
| int32_t bitmask_index = curr_token_id / 32; | |
| if (bitmask_index < 0 || bitmask_index >= bitmask_size) { | |
| // Out of bounds, skip this branch | |
| return; | |
| } | |
| accepted = (parent_bitmask[bitmask_index] & (1 << (curr_token_id % 32))) != 0; |
This PR adds a new TraverseDraftTree function that efficiently traverses a speculative decoding tree and generates grammar-aware token bitmasks for each node in a single pass.
Signed-off-by: Ubospica ubospica@gmail.com