Expose draft tree traversal on GrammarMatcher

[Ubospica]

Summary

• Move draft tree traversal into the public GrammarMatcher API for C++ and Python callers.
• Keep the Python testing wrapper as a backward-compatible shim while removing the C++ testing free function.
• Update speculative decoding tests to cover the member API and compatibility wrapper.

Validation

• cmake --build build -j2
• python -m pytest tests/python/test_speculative_decoding.py
• python -m ruff check python/xgrammar/matcher.py python/xgrammar/testing.py tests/python/test_speculative_decoding.py
• git diff --check

[gemini-code-assist]

Code Review

This pull request moves the TraverseDraftTree logic from the testing namespace into the core GrammarMatcher class, exposing it through C++ and Python APIs for speculative decoding. The implementation includes a recursive DFS traversal to compute token bitmasks across a draft tree. Feedback highlights critical issues in the traversal logic, including potential out-of-bounds access, incorrect state rollback, and the need for better bitmask initialization for rejected or terminated nodes. Additionally, suggestions were made to improve input tensor validation for device placement and dimensions.

[gemini-code-assist]

There are several issues in the recursive traversal logic:

1. OOB Access: current_token_id is used to index parent_bitmask without checking if it's within the valid range [0, vocab_size). This can lead to crashes or undefined behavior if the draft model produces an invalid token ID.
2. Incorrect Rollback: If matcher.AcceptToken fails (returns false), the code still calls matcher.Rollback(1) at the end of the block. This will incorrectly pop a state that was pushed by a previous successful token acceptance, corrupting the matcher's state for other branches of the tree.
3. Uninitialized Bitmasks: If a node is not accepted (either by the parent bitmask or the grammar), its corresponding row in token_bitmask is never filled. Since token_bitmask is typically initialized to all ones (unmasked) in Python, these rejected nodes will incorrectly appear as having no constraints.
4. Termination Handling: If the matcher is terminated, FillNextTokenBitmask cannot be called. The bitmask for such nodes should be zeroed out to indicate no further tokens are allowed.

  bool accepted;
  if (current_position == 0) {
    // The first token generated by the target model is always accepted.
    accepted = true;
  } else {
    int32_t current_token_id = static_cast<int32_t>(draft_tokens[current_position]);
    // Check if the token is within the valid range for the bitmask.
    if (current_token_id < 0 || current_token_id / 32 >= bitmask_size) {
      accepted = false;
    } else {
      int32_t* parent_bitmask = bitmask_data + parent_position * bitmask_size;
      // 32 boolean bitmask values are packed into 32-bit integers.
      accepted = (parent_bitmask[current_token_id / 32] & (1 << (current_token_id % 32))) != 0;
    }

    // Check timeout for non-root nodes so the root token mask is still computed.
    if (accepted && time_threshold > 0) {
      auto elapsed = std::chrono::duration<double>(Clock::now() - start_time).count();
      if (elapsed > time_threshold) {
        return false;
      }
    }
  }

  if (accepted) {
    bool token_accepted = true;
    if (current_position != 0) {
      token_accepted = matcher.AcceptToken(draft_tokens[current_position]);
    }

    if (token_accepted) {
      if (!matcher.IsTerminated()) {
        matcher.FillNextTokenBitmask(token_bitmask, current_position);

        if (retrieve_next_token[current_position] != -1) {
          bool success = TraverseDraftTreeRecursive(
              retrieve_next_token[current_position],
              current_position,
              retrieve_next_token,
              retrieve_next_sibling,
              draft_tokens,
              matcher,
              token_bitmask,
              time_threshold,
              start_time
          );
          if (!success) {
            if (current_position != 0) {
              matcher.Rollback(1);
            }
            return false;
          }
        }
      } else {
        // If terminated, no more tokens are allowed from this node.
        std::fill_n(bitmask_data + current_position * bitmask_size, bitmask_size, 0);
      }

      if (current_position != 0) {
        matcher.Rollback(1);
      }
    } else {
      // Token rejected by grammar: zero out the bitmask for this node.
      std::fill_n(bitmask_data + current_position * bitmask_size, bitmask_size, 0);
    }
  } else {
    // Node not accepted: zero out the bitmask for this node.
    std::fill_n(bitmask_data + current_position * bitmask_size, bitmask_size, 0);
  }

[gemini-code-assist]

The input tensors should be validated for device (must be CPU) and dimensions (ndim) before accessing their data pointers. Additionally, token_bitmask should be checked to ensure its batch size matches the number of nodes in the tree.

  XGRAMMAR_CHECK(retrieve_next_token->ndim == 1 && retrieve_next_token->dtype.code == kDLInt && retrieve_next_token->dtype.bits == 64)
      << "The retrieve_next_token tensor must be a 1D int64 tensor";
  XGRAMMAR_CHECK(retrieve_next_sibling->ndim == 1 && retrieve_next_sibling->dtype.code == kDLInt && retrieve_next_sibling->dtype.bits == 64)
      << "The retrieve_next_sibling tensor must be a 1D int64 tensor";
  XGRAMMAR_CHECK(draft_tokens->ndim == 1 && draft_tokens->dtype.code == kDLInt && draft_tokens->dtype.bits == 64)
      << "The draft_tokens tensor must be a 1D int64 tensor";
  XGRAMMAR_CHECK(token_bitmask->ndim == 2 && token_bitmask->dtype.code == kDLInt && token_bitmask->dtype.bits == 32)
      << "The token_bitmask tensor must be a 2D int32 tensor";

  auto check_cpu = [](const DLTensor* tensor, const char* name) {
    XGRAMMAR_CHECK(tensor->device.device_type == kDLCPU ||
                   tensor->device.device_type == kDLCUDAHost ||
                   tensor->device.device_type == kDLROCMHost)
        << "The " << name << " tensor must be on CPU";
  };
  check_cpu(retrieve_next_token, "retrieve_next_token");
  check_cpu(retrieve_next_sibling, "retrieve_next_sibling");
  check_cpu(draft_tokens, "draft_tokens");
  check_cpu(token_bitmask, "token_bitmask");

  XGRAMMAR_CHECK(retrieve_next_token->shape[0] == retrieve_next_sibling->shape[0])
      << "The retrieve_next_token and retrieve_next_sibling tensors must have the same length";
  XGRAMMAR_CHECK(retrieve_next_token->shape[0] == draft_tokens->shape[0])
      << "The retrieve_next_token and draft_tokens tensors must have the same length";
  XGRAMMAR_CHECK(retrieve_next_token->shape[0] == token_bitmask->shape[0])
      << "The token_bitmask batch size must match the number of nodes in the tree";

[Copilot]

Pull request overview

This PR promotes speculative-decoding draft-tree traversal from a C++ testing helper into the public GrammarMatcher API (C++ and Python), while keeping a Python compatibility shim and updating tests to cover both the new member API and the old wrapper.

Changes:

• Add GrammarMatcher::TraverseDraftTree(...) to the public C++ API and wire it through the TVM FFI.
• Expose GrammarMatcher.traverse_draft_tree(...) in Python and keep xgrammar.testing._traverse_draft_tree as a backward-compatible wrapper.
• Update speculative decoding tests to use the new member API and to explicitly test the compatibility shim.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
cpp/grammar_matcher.cc	Moves the DFS traversal implementation into GrammarMatcher as a public member.
include/xgrammar/matcher.h	Adds the public C++ declaration and docs for TraverseDraftTree.
cpp/tvm_ffi/tvm_ffi.cc	Exposes traverse_draft_tree as an FFI method and updates the legacy testing binding to call the member API.
cpp/testing.h	Removes the C++ testing free-function declaration for draft-tree traversal.
cpp/testing.cc	Removes the C++ testing free-function implementation (leaving other testing helpers).
python/xgrammar/matcher.py	Adds the public Python GrammarMatcher.traverse_draft_tree method.
python/xgrammar/testing.py	Converts _traverse_draft_tree into a compatibility wrapper calling the new member method.
tests/python/test_speculative_decoding.py	Updates tests to call the member API and adds coverage for the compatibility wrapper.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

GrammarMatcher::TraverseDraftTree validates dtypes and matching lengths, but does not validate that the tensors are CPU-accessible and have the expected ranks/shapes (e.g., token_bitmask is assumed 2D and accessed via shape[1] in the recursion). Passing a 1D bitmask or any GPU-only tensor can lead to out-of-bounds reads or dereferencing device memory in host code. Please add explicit checks for ndim (retrieve_* and draft_tokens are 1D; token_bitmask is 2D), token_bitmask.shape[0] == num_nodes, and that all inputs are on CPU/host-accessible devices before starting the traversal.

[Copilot]

TraverseDraftTreeRecursive computes parent_bitmask = bitmask_data + parent_position * bitmask_size for every non-root node without validating that parent_position is non-negative. If the input tree is malformed (e.g., retrieve_next_sibling[0] != -1 or any traversal path reaches a node with parent_position == -1), this will index before the buffer and cause memory corruption. Please enforce the invariant (no root siblings / valid parent_position for non-root nodes) with a check, or restructure the logic to avoid using an invalid parent_position.

[Copilot]

This new public Python API method routes through _core.testing._traverse_draft_tree, which couples the stable GrammarMatcher API to the internal testing FFI surface. Since the C++/FFI now exposes GrammarMatcher.traverse_draft_tree, prefer calling self._handle.traverse_draft_tree(...) directly here and reserve _core.testing.* for the backward-compat shim in xgrammar.testing.

Suggested change

	return _core.testing._traverse_draft_tree(
	retrieve_next_token,
	retrieve_next_sibling,
	draft_tokens,
	self._handle,
	return self._handle.traverse_draft_tree(
	retrieve_next_token,
	retrieve_next_sibling,
	draft_tokens,

[Ubospica]

/gemini review

[gemini-code-assist]

Code Review

This pull request moves the TraverseDraftTree functionality from a testing utility into the core GrammarMatcher class, making it part of the public API. The implementation includes a recursive DFS traversal of the speculative decoding tree to generate token bitmasks. Feedback focuses on memory safety and robustness: specifically, the current implementation ignores the byte_offset and strides fields of the DLTensor objects, which could lead to incorrect memory access if tensors are sliced or padded. Additionally, it is recommended to replace the recursive sibling traversal with an iterative loop to prevent potential stack overflow issues on wide trees.

[gemini-code-assist]

The implementation ignores the byte_offset of the token_bitmask DLTensor. According to the DLPack specification, the actual data pointer should be calculated as (char*)data + byte_offset. Additionally, the code assumes the tensor is contiguous by using shape[1] as the row stride. If the input tensor is a slice or has padding, this will lead to incorrect memory access. It is recommended to use strides[0] (if not null) to navigate between rows.

[gemini-code-assist]

Directly accessing retrieve_next_sibling->data ignores the byte_offset field of the DLTensor. This can lead to incorrect results if the tensor is a slice. Please use static_cast<const char*>(tensor->data) + tensor->byte_offset to calculate the correct pointer. This also applies to the pointers passed to TraverseDraftTreeRecursive on lines 1320-1322.

[gemini-code-assist]

The sibling traversal is implemented recursively. While this is tail-recursive, it is safer and more idiomatic to use an iterative loop for siblings to avoid potential stack overflow issues on very wide trees, especially since this is now part of the core public API. A while loop similar to the one in ClearDraftTreeSiblingChain would be more robust.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The timeout check is gated on accepted, so large trees containing many rejected nodes can effectively ignore time_threshold and run far longer than the caller requested. Consider checking the elapsed time for every non-root node (as the previous implementation did), or at least for every visited node before traversing siblings, so time_threshold bounds total traversal time regardless of acceptance.

[Copilot]

TraverseDraftTree dereferences raw DLTensor->data pointers later in the function, but currently only checks retrieve_next_sibling->data for null. Please add explicit non-null (and non-empty where relevant) checks for retrieve_next_token->data, draft_tokens->data, and token_bitmask->data as well, to avoid potential null dereference crashes when called through DLPack/FFI.

[Copilot]

TraverseDraftTreeRecursive uses recursion for both child and sibling traversal. Since this is now part of the public GrammarMatcher API, very deep/degenerate draft trees (e.g., long linear chains with timeout disabled) can cause stack overflow. Consider rewriting this traversal iteratively with an explicit stack (or enforcing a maximum depth) to make the API robust for untrusted/large inputs.

Suggested change

	int32_t* bitmask_data = reinterpret_cast<int32_t*>(token_bitmask->data);
	int32_t bitmask_size = static_cast<int32_t>(token_bitmask->shape[1]);
	bool accepted;
	if (current_position == 0) {
	// The first token generated by the target model is always accepted.
	accepted = true;
	} else {
	XGRAMMAR_CHECK(parent_position >= 0)
	<< "Non-root draft tree nodes must have a valid parent position";
	int64_t current_token_id = draft_tokens[current_position];
	if (current_token_id < 0 || current_token_id >= static_cast<int64_t>(bitmask_size) * 32) {
	accepted = false;
	} else {
	int32_t* parent_bitmask = bitmask_data + parent_position * bitmask_size;
	// 32 boolean bitmask values are packed into 32-bit integers.
	uint32_t token_mask = uint32_t{1} << static_cast<uint32_t>(current_token_id % 32);
	accepted = (static_cast<uint32_t>(parent_bitmask[current_token_id / 32]) & token_mask) != 0;
	}
	// Check timeout for non-root nodes so the root token mask is still computed.
	if (accepted && time_threshold > 0) {
	auto elapsed = std::chrono::duration<double>(Clock::now() - start_time).count();
	if (elapsed > time_threshold) {
	return false;
	}
	}
	}
	if (accepted) {
	bool token_accepted = true;
	if (current_position != 0) {
	token_accepted = matcher.AcceptToken(static_cast<int32_t>(draft_tokens[current_position]));
	}
	if (token_accepted) {
	if (!matcher.IsTerminated()) {
	matcher.FillNextTokenBitmask(token_bitmask, current_position);
	if (retrieve_next_token[current_position] != -1) {
	bool success = TraverseDraftTreeRecursive(
	retrieve_next_token[current_position],
	current_position,
	retrieve_next_token,
	retrieve_next_sibling,
	draft_tokens,
	matcher,
	token_bitmask,
	time_threshold,
	start_time
	);
	if (!success) {
	if (current_position != 0) {
	matcher.Rollback(1);
	}
	return false;
	}
	}
	} else {
	ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position);
	}
	if (current_position != 0) {
	matcher.Rollback(1);
	}
	} else {
	ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position);
	}
	} else {
	ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position);
	}
	if (retrieve_next_sibling[current_position] != -1) {
	bool success = TraverseDraftTreeRecursive(
	retrieve_next_sibling[current_position],
	parent_position,
	retrieve_next_token,
	retrieve_next_sibling,
	draft_tokens,
	matcher,
	token_bitmask,
	time_threshold,
	start_time
	);
	if (!success) {
	return false;
	}
	struct TraversalFrame {
	int32_t current_position;
	int32_t parent_position;
	bool entered;
	bool rollback_on_exit;
	};
	int32_t* bitmask_data = reinterpret_cast<int32_t*>(token_bitmask->data);
	int32_t bitmask_size = static_cast<int32_t>(token_bitmask->shape[1]);
	std::vector<TraversalFrame> stack;
	stack.push_back({current_position, parent_position, false, false});
	while (!stack.empty()) {
	TraversalFrame& frame = stack.back();
	if (!frame.entered) {
	frame.entered = true;
	bool accepted;
	if (frame.current_position == 0) {
	// The first token generated by the target model is always accepted.
	accepted = true;
	} else {
	XGRAMMAR_CHECK(frame.parent_position >= 0)
	<< "Non-root draft tree nodes must have a valid parent position";
	int64_t current_token_id = draft_tokens[frame.current_position];
	if (current_token_id < 0 || current_token_id >= static_cast<int64_t>(bitmask_size) * 32) {
	accepted = false;
	} else {
	int32_t* parent_bitmask = bitmask_data + frame.parent_position * bitmask_size;
	// 32 boolean bitmask values are packed into 32-bit integers.
	uint32_t token_mask = uint32_t{1} << static_cast<uint32_t>(current_token_id % 32);
	accepted =
	(static_cast<uint32_t>(parent_bitmask[current_token_id / 32]) & token_mask) != 0;
	}
	// Check timeout for non-root nodes so the root token mask is still computed.
	if (accepted && time_threshold > 0) {
	auto elapsed = std::chrono::duration<double>(Clock::now() - start_time).count();
	if (elapsed > time_threshold) {
	return false;
	}
	}
	}
	if (!accepted) {
	ClearTokenBitmaskRow(bitmask_data, bitmask_size, frame.current_position);
	int64_t sibling_position = retrieve_next_sibling[frame.current_position];
	stack.pop_back();
	if (sibling_position != -1) {
	stack.push_back(
	{static_cast<int32_t>(sibling_position), frame.parent_position, false, false}
	);
	}
	continue;
	}
	bool token_accepted = true;
	if (frame.current_position != 0) {
	token_accepted = matcher.AcceptToken(static_cast<int32_t>(draft_tokens[frame.current_position]));
	}
	if (!token_accepted) {
	ClearTokenBitmaskRow(bitmask_data, bitmask_size, frame.current_position);
	int64_t sibling_position = retrieve_next_sibling[frame.current_position];
	stack.pop_back();
	if (sibling_position != -1) {
	stack.push_back(
	{static_cast<int32_t>(sibling_position), frame.parent_position, false, false}
	);
	}
	continue;
	}
	frame.rollback_on_exit = (frame.current_position != 0);
	if (matcher.IsTerminated()) {
	ClearTokenBitmaskRow(bitmask_data, bitmask_size, frame.current_position);
	} else {
	matcher.FillNextTokenBitmask(token_bitmask, frame.current_position);
	}
	int64_t child_position = retrieve_next_token[frame.current_position];
	if (child_position != -1 && !matcher.IsTerminated()) {
	stack.push_back(
	{static_cast<int32_t>(child_position), frame.current_position, false, false}
	);
	}
	continue;
	}
	int64_t sibling_position = retrieve_next_sibling[frame.current_position];
	int32_t sibling_parent_position = frame.parent_position;
	bool rollback_on_exit = frame.rollback_on_exit;
	stack.pop_back();
	if (rollback_on_exit) {
	matcher.Rollback(1);
	}
	if (sibling_position != -1) {
	stack.push_back(
	{static_cast<int32_t>(sibling_position), sibling_parent_position, false, false}
	);
	}