From 952b3caf18317b91dc072ebfa1c416a5cfd5eedb Mon Sep 17 00:00:00 2001 From: Linzhang Li Date: Mon, 4 May 2026 17:30:28 +0800 Subject: [PATCH] feat: use structural tags to enable strict tool calling and reasoning for more models (#21722) Signed-off-by: Yuchuan Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Co-authored-by: Ubospica Co-authored-by: Xinyuan Tong --- .github/workflows/nightly-test-npu.yml | 10 +- 3rdparty/amd/wheel/sglang/pyproject.toml | 2 +- python/pyproject.toml | 2 +- python/pyproject_cpu.toml | 2 +- python/pyproject_npu.toml | 2 +- python/pyproject_other.toml | 2 +- python/pyproject_xpu.toml | 1 + .../srt/entrypoints/openai/serving_chat.py | 8 + .../srt/function_call/base_format_detector.py | 52 +- .../srt/function_call/deepseekv32_detector.py | 3 + .../srt/function_call/deepseekv4_detector.py | 67 ++ .../srt/function_call/function_call_parser.py | 59 +- .../srt/function_call/gpt_oss_detector.py | 3 + .../srt/function_call/kimik2_detector.py | 10 + .../srt/function_call/qwen3_coder_detector.py | 5 +- python/sglang/srt/speculative/spec_utils.py | 6 +- .../test_function_call_parser.py | 666 +++++++++++++++++- .../spec/test_spec_utils_traverse_tree.py | 71 ++ 18 files changed, 922 insertions(+), 49 deletions(-) create mode 100644 python/sglang/srt/function_call/deepseekv4_detector.py create mode 100644 test/registered/unit/spec/test_spec_utils_traverse_tree.py diff --git a/.github/workflows/nightly-test-npu.yml b/.github/workflows/nightly-test-npu.yml index 205dcc02e..44071afc7 100644 --- a/.github/workflows/nightly-test-npu.yml +++ b/.github/workflows/nightly-test-npu.yml @@ -131,7 +131,7 @@ jobs: pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 - pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.32 numpy==1.26.4 dotenv + pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.2.0 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & @@ -197,7 +197,7 @@ jobs: pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 - pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.32 numpy==1.26.4 dotenv + pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.2.0 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & @@ -264,7 +264,7 @@ jobs: pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 - pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.32 numpy==1.26.4 dotenv + pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.2.0 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & @@ -331,7 +331,7 @@ jobs: pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 - pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.32 numpy==1.26.4 dotenv + pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.2.0 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & @@ -398,7 +398,7 @@ jobs: pip install sentence_transformers torchaudio==2.8.0 pip install protobuf==6.31.1 zss pre-commit wandb>=0.16.0 tenacity==8.3.0 loguru openpyxl latex2sympy2 zstandard transformers-stream-generator tqdm-multiprocess pycocoevalcap pip install yt-dlp sentencepiece==0.1.99 nltk av ftfy sqlitedict==2.1.0 sacrebleu>=1.5.0 pytablewriter black==24.1.0 isort==5.13.2 peft>=0.2.0 accelerate>=0.29.1 - pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.1.32 numpy==1.26.4 dotenv + pip install jsonlines httpx==0.25.0 evaluate>=0.4.0 datasets==2.16.1 numexpr xgrammar==0.2.0 numpy==1.26.4 dotenv git clone --branch v0.3.3 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git cd ./lmms-eval nohup pip install . > lmmslog.txt 2>&1 & diff --git a/3rdparty/amd/wheel/sglang/pyproject.toml b/3rdparty/amd/wheel/sglang/pyproject.toml index df1f27f58..41a012a77 100644 --- a/3rdparty/amd/wheel/sglang/pyproject.toml +++ b/3rdparty/amd/wheel/sglang/pyproject.toml @@ -62,7 +62,7 @@ runtime_common = [ "transformers==4.57.1", "uvicorn", "uvloop", - "xgrammar==0.1.32", + "xgrammar==0.2.0", "smg-grpc-servicer>=0.5.0", ] diff --git a/python/pyproject.toml b/python/pyproject.toml index 7fce5f7b9..9da94e25f 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -76,7 +76,7 @@ dependencies = [ "uvicorn", "uvloop", "watchfiles", - "xgrammar==0.1.32", + "xgrammar==0.2.0", "smg-grpc-servicer>=0.5.0", "kernels", ] diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index 677490ce1..963247d29 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -66,7 +66,7 @@ dependencies = [ "triton==3.5.0", "uvicorn", "uvloop", - "xgrammar==0.1.32", + "xgrammar==0.2.0", "smg-grpc-servicer>=0.5.0", ] diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 6b463293b..473e9dc5c 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -62,7 +62,7 @@ dependencies = [ "transformers==5.6.0", "uvicorn", "uvloop", - "xgrammar==0.1.32", + "xgrammar==0.2.0", "smg-grpc-servicer>=0.5.0", ] diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 7eed0643d..35f923f7e 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -62,7 +62,7 @@ runtime_common = [ "transformers==5.6.0", "uvicorn", "uvloop", - "xgrammar==0.1.32", + "xgrammar==0.2.0", "smg-grpc-servicer>=0.5.0", ] diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index b95017919..af4acb9bd 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -66,6 +66,7 @@ dependencies = [ "transformers==5.6.0", "uvicorn", "uvloop", + # "xgrammar==0.2.0", xgrammar depends on CUDA PyTorch and Triton only "smg-grpc-servicer>=0.5.0", ] diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index bb8c12596..72ef2780a 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -434,6 +434,13 @@ class OpenAIServingChat(OpenAIServingBase): self._patch_mistral_skip_special_tokens(request) + thinking_mode = self._get_reasoning_from_request(request) + # SGLang's ReasonerGrammarBackend owns the reasoning prefix + # when --reasoning-parser is configured, so builtin xgrammar + # tags must describe only the post-reasoning tool-call suffix. + xgrammar_reasoning = thinking_mode and ( + self.tokenizer_manager.server_args.reasoning_parser is not None + ) tool_call_constraint = None # Apply chat template and its stop strings @@ -453,6 +460,7 @@ class OpenAIServingChat(OpenAIServingBase): tool_call_constraint = parser.get_structure_constraint( request.tool_choice, parallel_tool_calls=request.parallel_tool_calls, + thinking_mode=xgrammar_reasoning, ) # Fallback: use generic JSON schema for required/named tool choice # only when no parser-specific constraint was set diff --git a/python/sglang/srt/function_call/base_format_detector.py b/python/sglang/srt/function_call/base_format_detector.py index b89d51796..726288ea3 100644 --- a/python/sglang/srt/function_call/base_format_detector.py +++ b/python/sglang/srt/function_call/base_format_detector.py @@ -1,13 +1,19 @@ import json import logging from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import Any, Dict, List, Literal, Optional, Union import orjson from partial_json_parser.core.exceptions import MalformedJSON from partial_json_parser.core.options import Allow -from sglang.srt.entrypoints.openai.protocol import Tool +try: + from xgrammar import StructuralTag, get_model_structural_tag +except ImportError: + StructuralTag = Any + get_model_structural_tag = None + +from sglang.srt.entrypoints.openai.protocol import Tool, ToolChoice from sglang.srt.environ import envs from sglang.srt.function_call.core_types import ( StreamingParseResult, @@ -361,3 +367,45 @@ class BaseFormatDetector(ABC): A function that takes a tool name (str) and returns StructureInfo """ raise NotImplementedError() + + def get_structural_tag_name(self) -> Optional[str]: + """Return the XGrammar model name for native structural tags, if supported.""" + return None + + def get_structural_tag( + self, + tools: Union[List[Tool], None] = None, + tool_choice: Union[ToolChoice, Literal["auto", "required"]] = "auto", + thinking_mode: bool = False, + ) -> Optional[StructuralTag]: + """ + Return a model-native XGrammar structural tag when supported. + + Args: + tools: List of available tools + tool_choice: The tool choice setting from the request + thinking_mode: Whether to include the model's reasoning prefix in + the returned structural tag. Pass False when SGLang's + ReasonerGrammarBackend will own the ... prefix + (the typical case when --reasoning-parser is configured) so + only one layer constrains the reasoning section. + + Returns: + StructuralTag if this detector supports model-native tags, otherwise None + """ + structural_tag_name = self.get_structural_tag_name() + if not structural_tag_name or get_model_structural_tag is None: + return None + + converted_tools = [tool.model_dump() for tool in tools or []] + converted_tool_choice = ( + tool_choice.model_dump() + if isinstance(tool_choice, ToolChoice) + else tool_choice + ) + return get_model_structural_tag( + model=structural_tag_name, + tools=converted_tools, + tool_choice=converted_tool_choice, + reasoning=thinking_mode, + ) diff --git a/python/sglang/srt/function_call/deepseekv32_detector.py b/python/sglang/srt/function_call/deepseekv32_detector.py index d9228743b..5bc3fdcb2 100644 --- a/python/sglang/srt/function_call/deepseekv32_detector.py +++ b/python/sglang/srt/function_call/deepseekv32_detector.py @@ -351,3 +351,6 @@ class DeepSeekV32Detector(BaseFormatDetector): end="", trigger="<｜DSML｜invoke", ) + + def get_structural_tag_name(self) -> str: + return "deepseek_v3_2" diff --git a/python/sglang/srt/function_call/deepseekv4_detector.py b/python/sglang/srt/function_call/deepseekv4_detector.py new file mode 100644 index 000000000..10d4d559f --- /dev/null +++ b/python/sglang/srt/function_call/deepseekv4_detector.py @@ -0,0 +1,67 @@ +import logging + +from sglang.srt.function_call.deepseekv32_detector import DeepSeekV32Detector + +logger = logging.getLogger(__name__) + + +class DeepSeekV4Detector(DeepSeekV32Detector): + """ + Detector for DeepSeek V4 model function call format. + + The DeepSeek V4 format uses XML-like DSML tags to delimit function calls. + Supports two parameter formats: + + Format 1 - XML Parameter Tags: + ``` + <｜DSML｜tool_calls> + <｜DSML｜invoke name="function_name"> + <｜DSML｜parameter name="param_name" string="true">value + ... + + + ``` + + Format 2 - Direct JSON: + ``` + <｜DSML｜tool_calls> + <｜DSML｜invoke name="function_name"> + { + "param_name": "value" + } + + + ``` + + Examples: + ``` + <｜DSML｜tool_calls> + <｜DSML｜invoke name="get_favorite_tourist_spot"> + <｜DSML｜parameter name="city" string="true">San Francisco + + + + <｜DSML｜tool_calls> + <｜DSML｜invoke name="get_favorite_tourist_spot"> + { "city": "San Francisco" } + + + ``` + + Key Components: + - Tool Calls Section: Wrapped between `<｜DSML｜tool_calls>` and `` + - Individual Tool Call: Wrapped between `<｜DSML｜invoke name="...">` and `` + - Parameters: Either XML tags or direct JSON format + - Supports multiple tool calls + + Reference: DeepSeek V4 format specification + """ + + def __init__(self): + super().__init__() + self.bot_token = "<｜DSML｜tool_calls>" + self.eot_token = "" + self.function_calls_regex = r"<｜DSML｜tool_calls>(.*?)" + + def get_structural_tag_name(self) -> str: + return "deepseek_v4" diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 08585e556..c9524ac5a 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -152,7 +152,7 @@ class FunctionCallParser: return final_normal_text, final_calls - def get_structure_tag( + def get_legacy_structural_tag( self, at_least_one: bool = False ) -> StructuralTagResponseFormat: """ @@ -208,6 +208,7 @@ class FunctionCallParser: self, tool_choice: Union[ToolChoice, Literal["auto", "required"]], parallel_tool_calls: bool = True, + thinking_mode: bool = False, ) -> Optional[ToolCallConstraint]: """ Returns the appropriate structure constraint for tool calls based on the tool_choice. @@ -220,27 +221,37 @@ class FunctionCallParser: A tuple of (constraint_type, constraint_value) to be added to sampling parameters, or None if no constraint applies. """ - # NOTE: structural_tag only supports JSON-compatible content between the begin and end. - # It cannot parse or validate function call Pythonic or XML-ish syntax. - if self.detector.supports_structural_tag(): - # For "required"/named: always use structural_tag to preserve the - # model's native tool call format. Schema is only included when - # strict=True, per OpenAI protocol semantics. - # For "auto": only constrain when strict is enabled. - is_required = tool_choice == "required" or isinstance( - tool_choice, ToolChoice - ) - if is_required or ( - tool_choice == "auto" - and ( - any(tool.function.strict for tool in self.tools) - or self.tool_strict_level >= ToolStrictLevel.FUNCTION + is_required = tool_choice == "required" or isinstance(tool_choice, ToolChoice) + should_constrain_auto = tool_choice == "auto" and ( + any(tool.function.strict for tool in self.tools) + or self.tool_strict_level >= ToolStrictLevel.FUNCTION + ) + + # Highest priority: model-native structural_tag when available. + try: + if is_required or should_constrain_auto: + structural_tag = self.detector.get_structural_tag( + tools=self.tools, + thinking_mode=thinking_mode, + tool_choice=tool_choice, ) - ): - tag = self.get_structure_tag(at_least_one=is_required) - return ("structural_tag", tag) - elif tool_choice == "required" or isinstance(tool_choice, ToolChoice): - json_schema = get_json_schema_constraint( - self.tools, tool_choice, parallel_tool_calls=parallel_tool_calls - ) - return ("json_schema", json_schema) + if structural_tag is not None: + return ("structural_tag", structural_tag) + + # Fallback to legacy structural tag if model-native tag is not supported. + if self.detector.supports_structural_tag(): + # For "required"/named: always use structural_tag to preserve the + # model's native tool call format. Schema is only included when + # strict=True, per OpenAI protocol semantics. + # For "auto": only constrain when strict is enabled. + tag = self.get_legacy_structural_tag(at_least_one=is_required) + return ("structural_tag", tag) + + if tool_choice == "required" or isinstance(tool_choice, ToolChoice): + json_schema = get_json_schema_constraint( + self.tools, tool_choice, parallel_tool_calls=parallel_tool_calls + ) + return ("json_schema", json_schema) + except Exception as e: + logger.error(f"Error getting structure constraint: {e}") + return None diff --git a/python/sglang/srt/function_call/gpt_oss_detector.py b/python/sglang/srt/function_call/gpt_oss_detector.py index b3fd5ac61..b7234262f 100644 --- a/python/sglang/srt/function_call/gpt_oss_detector.py +++ b/python/sglang/srt/function_call/gpt_oss_detector.py @@ -239,3 +239,6 @@ class GptOssDetector(BaseFormatDetector): def structure_info(self) -> _GetInfoFunc: raise NotImplementedError("structure_info not used with HarmonyParser") + + def get_structural_tag_name(self) -> str: + return "harmony" diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py index 21cf46cb0..113031c8e 100644 --- a/python/sglang/srt/function_call/kimik2_detector.py +++ b/python/sglang/srt/function_call/kimik2_detector.py @@ -253,3 +253,13 @@ class KimiK2Detector(BaseFormatDetector): ) return get_info + + # Kimi stays on the SGLang legacy structural tag path. xgrammar 0.2.0's + # get_kimi_structural_tag(tool_choice="auto") emits a bare + # <|tool_call_begin|>...<|tool_call_end|> grammar without the + # <|tool_calls_section_begin|>/<|tool_calls_section_end|> wrapper Kimi's + # chat template uses, and KimiK2Detector.has_tool_call() keys off the + # section marker — bare tool calls would be silently dropped. Inheriting + # the base get_structural_tag_name (returns None) keeps FunctionCallParser + # on the legacy path, whose structure_info bakes the section markers in. + # TODO: re-enable the builtin once https://github.com/mlc-ai/xgrammar/issues/622 is fixed. diff --git a/python/sglang/srt/function_call/qwen3_coder_detector.py b/python/sglang/srt/function_call/qwen3_coder_detector.py index 9dd77903d..025404572 100644 --- a/python/sglang/srt/function_call/qwen3_coder_detector.py +++ b/python/sglang/srt/function_call/qwen3_coder_detector.py @@ -468,7 +468,10 @@ class Qwen3CoderDetector(BaseFormatDetector): return StreamingParseResult(calls=calls, normal_text=normal_text) def supports_structural_tag(self) -> bool: - return False + return True def structure_info(self) -> _GetInfoFunc: raise NotImplementedError + + def get_structural_tag_name(self) -> str: + return "qwen_3_coder" diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 46f6c00b1..aebdc4fb4 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -612,14 +612,14 @@ def traverse_tree( if accepted: if curr != 0: # Accept the current token - grammar.accept_token(draft_tokens[curr]) + grammar.accept_token(int(draft_tokens[curr])) if not grammar.is_terminated(): # Generate the bitmask for the current token grammar.fill_vocab_mask(allocate_token_bitmask, curr) if retrieve_next_token[curr] != -1: # Visit the child node dfs( - retrieve_next_token[curr], + int(retrieve_next_token[curr]), retrieve_next_token, retrieve_next_sibling, curr, @@ -632,7 +632,7 @@ def traverse_tree( if retrieve_next_sibling[curr] != -1: # Visit the sibling node dfs( - retrieve_next_sibling[curr], + int(retrieve_next_sibling[curr]), retrieve_next_token, retrieve_next_sibling, parent_pos, diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py index d021694cb..3fb68066e 100644 --- a/test/registered/unit/function_call/test_function_call_parser.py +++ b/test/registered/unit/function_call/test_function_call_parser.py @@ -1,10 +1,16 @@ import json import unittest -from sglang.srt.entrypoints.openai.protocol import Function, Tool +from sglang.srt.entrypoints.openai.protocol import ( + Function, + Tool, + ToolChoice, + ToolChoiceFuncName, +) from sglang.srt.function_call.base_format_detector import BaseFormatDetector from sglang.srt.function_call.core_types import StreamingParseResult from sglang.srt.function_call.deepseekv3_detector import DeepSeekV3Detector +from sglang.srt.function_call.deepseekv4_detector import DeepSeekV4Detector from sglang.srt.function_call.deepseekv32_detector import DeepSeekV32Detector from sglang.srt.function_call.gemma4_detector import ( Gemma4Detector, @@ -15,6 +21,7 @@ from sglang.srt.function_call.gemma4_detector import ( from sglang.srt.function_call.gigachat3_detector import GigaChat3Detector from sglang.srt.function_call.glm4_moe_detector import Glm4MoeDetector from sglang.srt.function_call.glm47_moe_detector import Glm47MoeDetector +from sglang.srt.function_call.gpt_oss_detector import GptOssDetector from sglang.srt.function_call.json_array_parser import JsonArrayParser from sglang.srt.function_call.kimik2_detector import KimiK2Detector from sglang.srt.function_call.lfm2_detector import Lfm2Detector @@ -1632,6 +1639,478 @@ class TestDeepSeekV32Detector(unittest.TestCase): params = json.loads(tool_calls_by_index[0]["parameters"]) self.assertEqual(params, {}) + def test_get_model_structural_tag(self): + import xgrammar as xgr + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + tool_choice_name = ToolChoiceFuncName(name="search") + tool_choice = ToolChoice(function=tool_choice_name) + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + +class TestDeepSeekV4Detector(unittest.TestCase): + def setUp(self): + """Set up test tools and detector for DeepSeekV4 format testing.""" + self.tools = [ + Tool( + type="function", + function=Function( + name="search", + description="Searches for information related to query and displays topn results.", + parameters={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query string", + }, + "topn": { + "type": "integer", + "description": "Number of top results to display", + "default": 10, + }, + "source": { + "type": "string", + "description": "Source to search within", + "enum": ["web", "news"], + "default": "web", + }, + }, + "required": ["query"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="get_favorite_tourist_spot", + description="Return the favorite tourist spot for a given city.", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + ), + ), + ] + self.detector = DeepSeekV4Detector() + from sglang.srt.utils.hf_transformers_utils import get_tokenizer + + self.tokenizer = get_tokenizer("deepseek-ai/DeepSeek-V3.2") + self.interval = 1 + + def test_detect_and_parse_xml_format(self): + """Test parsing standard XML format (DSML)""" + text = """I'll help you with information about San Francisco and get its favorite tourist spot for you.\n\n + <｜DSML｜tool_calls>\n + <｜DSML｜invoke name="get_favorite_tourist_spot">\n + <｜DSML｜parameter name="city" string="true">San Francisco\n + \n + <｜DSML｜invoke name="search"> + <｜DSML｜parameter name="query" string="true">WebNav benchmark + <｜DSML｜parameter name="topn" string="false">10 + <｜DSML｜parameter name="source" string="true">web + + + """ + result = self.detector.detect_and_parse(text, self.tools) + + self.assertIn("I'll help you with information", result.normal_text) + self.assertEqual(len(result.calls), 2) + + # Check first call + call1 = result.calls[0] + self.assertEqual(call1.name, "get_favorite_tourist_spot") + params1 = json.loads(call1.parameters) + self.assertEqual(params1["city"], "San Francisco") + + # Check second call + call2 = result.calls[1] + self.assertEqual(call2.name, "search") + params2 = json.loads(call2.parameters) + self.assertEqual(params2["query"], "WebNav benchmark") + self.assertEqual(params2["topn"], 10) + self.assertEqual(params2["source"], "web") + + def test_detect_and_parse_json_format(self): + """Test parsing JSON format inside invoke tags""" + text = """I'll help you with information about San Francisco and get its favorite tourist spot for you. + + <｜DSML｜tool_calls> + <｜DSML｜invoke name="get_favorite_tourist_spot"> + { + "city": "San Francisco" + } + + <｜DSML｜invoke name="search"> + { + "query": "WebNav benchmark", + "topn": 10, + "source": "web" + } + + + """ + result = self.detector.detect_and_parse(text, self.tools) + + self.assertIn("I'll help you with information", result.normal_text) + self.assertEqual(len(result.calls), 2) + + # Check first call + call1 = result.calls[0] + self.assertEqual(call1.name, "get_favorite_tourist_spot") + params1 = json.loads(call1.parameters) + self.assertEqual(params1["city"], "San Francisco") + + # Check second call + call2 = result.calls[1] + self.assertEqual(call2.name, "search") + params2 = json.loads(call2.parameters) + self.assertEqual(params2["query"], "WebNav benchmark") + self.assertEqual(params2["topn"], 10) + self.assertEqual(params2["source"], "web") + + def test_streaming_xml_format(self): + """Test streaming parsing of XML format""" + text = """<｜DSML｜tool_calls> + <｜DSML｜invoke name="get_favorite_tourist_spot"> + <｜DSML｜parameter name="city" string="true">San Francisco + <｜DSML｜parameter name="another_city" string="true">London + <｜DSML｜parameter name="topn" string="false">10 + <｜DSML｜parameter name="obj" string="false">{"name": "John", "age": 30} + + """ + + input_ids = self.tokenizer.encode(text, add_special_tokens=False) + chunk_ids = [ + input_ids[i : i + self.interval] + for i in range(0, len(input_ids), self.interval) + ] + chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] + + tool_calls_by_index = {} + + num_tool_call_chunks = 0 + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + for call in result.calls: + num_tool_call_chunks += 1 + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + self.assertGreater(num_tool_call_chunks, 8) + + self.assertEqual(len(tool_calls_by_index), 1) + self.assertEqual(tool_calls_by_index[0]["name"], "get_favorite_tourist_spot") + params = json.loads(tool_calls_by_index[0]["parameters"]) + self.assertEqual(params["city"], "San Francisco") + self.assertEqual(params["another_city"], "London") + self.assertEqual(params["topn"], 10) + self.assertEqual(params["obj"]["name"], "John") + self.assertEqual(params["obj"]["age"], 30) + + def test_streaming_json_format(self): + """Test streaming parsing of JSON format""" + text = """<｜DSML｜tool_calls> + <｜DSML｜invoke name="get_favorite_tourist_spot"> + { + "city": "San Francisco", + "another_city": "London", + "topn": 10, + "obj": { + "name": "John", + "age": 30 + } + } + + """ + + input_ids = self.tokenizer.encode(text, add_special_tokens=False) + chunk_ids = [ + input_ids[i : i + self.interval] + for i in range(0, len(input_ids), self.interval) + ] + chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] + + tool_calls_by_index = {} + + num_tool_call_chunks = 0 + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + for call in result.calls: + num_tool_call_chunks += 1 + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + self.assertGreater(num_tool_call_chunks, 8) + self.assertEqual(len(tool_calls_by_index), 1) + self.assertEqual(tool_calls_by_index[0]["name"], "get_favorite_tourist_spot") + + # Clean up parameters string if needed (trim whitespace) + params_str = tool_calls_by_index[0]["parameters"].strip() + params = json.loads(params_str) + self.assertEqual(params["city"], "San Francisco") + + def test_detect_and_parse_no_parameters(self): + """Test parsing function calls with no parameters (non-streaming)""" + # Add a no-parameter tool + tools_with_no_param = self.tools + [ + Tool( + type="function", + function=Function( + name="get_date", + description="Get the current date.", + parameters={"type": "object", "properties": {}}, + ), + ), + ] + + text = """Let me get the current date for you. + +<｜DSML｜tool_calls> +<｜DSML｜invoke name="get_date"> + +""" + + result = self.detector.detect_and_parse(text, tools_with_no_param) + + self.assertIn("Let me get the current date", result.normal_text) + self.assertEqual(len(result.calls), 1) + + call = result.calls[0] + self.assertEqual(call.name, "get_date") + params = json.loads(call.parameters) + self.assertEqual(params, {}) + + def test_streaming_no_parameters(self): + """Test streaming parsing of function calls with no parameters. + + This test verifies the fix for the bug where functions with no parameters + were being silently skipped in streaming mode. + """ + # Add a no-parameter tool + tools_with_no_param = self.tools + [ + Tool( + type="function", + function=Function( + name="get_date", + description="Get the current date.", + parameters={"type": "object", "properties": {}}, + ), + ), + ] + + text = """<｜DSML｜tool_calls> +<｜DSML｜invoke name="get_date"> + +""" + + # Reset detector state + self.detector = DeepSeekV4Detector() + + # Simulate streaming by splitting into small chunks + input_ids = self.tokenizer.encode(text, add_special_tokens=False) + chunk_ids = [ + input_ids[i : i + self.interval] + for i in range(0, len(input_ids), self.interval) + ] + chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] + + tool_calls_by_index = {} + + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, tools_with_no_param) + for call in result.calls: + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + # Verify that the no-parameter function was correctly parsed + self.assertEqual( + len(tool_calls_by_index), 1, "Should have exactly one tool call" + ) + self.assertEqual(tool_calls_by_index[0]["name"], "get_date") + + # Parameters should be empty JSON object + params_str = tool_calls_by_index[0]["parameters"].strip() + params = json.loads(params_str) + self.assertEqual(params, {}) + + def test_streaming_no_parameters_with_whitespace(self): + """Test streaming parsing when invoke content has only whitespace (newlines).""" + tools_with_no_param = self.tools + [ + Tool( + type="function", + function=Function( + name="get_date", + description="Get the current date.", + parameters={"type": "object", "properties": {}}, + ), + ), + ] + + # This format has newlines inside the invoke tag (common model output) + text = """<｜DSML｜tool_calls> +<｜DSML｜invoke name="get_date"> + + +""" + + # Reset detector state + self.detector = DeepSeekV4Detector() + + input_ids = self.tokenizer.encode(text, add_special_tokens=False) + chunk_ids = [ + input_ids[i : i + self.interval] + for i in range(0, len(input_ids), self.interval) + ] + chunks = [self.tokenizer.decode(chunk_id) for chunk_id in chunk_ids] + + tool_calls_by_index = {} + + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, tools_with_no_param) + for call in result.calls: + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + # Should still parse correctly even with whitespace-only content + self.assertEqual( + len(tool_calls_by_index), 1, "Should have exactly one tool call" + ) + self.assertEqual(tool_calls_by_index[0]["name"], "get_date") + params = json.loads(tool_calls_by_index[0]["parameters"]) + self.assertEqual(params, {}) + + def test_get_model_structural_tag(self): + import xgrammar as xgr + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + tool_choice_name = ToolChoiceFuncName(name="search") + tool_choice = ToolChoice(function=tool_choice_name) + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + class TestQwen3CoderDetector(unittest.TestCase): """Test suite for Qwen3CoderDetector.""" @@ -1985,6 +2464,148 @@ class TestQwen3CoderDetector(unittest.TestCase): self.assertFalse(self.detector.has_tool_call("plain text only")) self.assertFalse(self.detector.has_tool_call("")) + # ==================== Structural tag (xgrammar builtin) ==================== + # Qwen3 Coder uses the new builtin structural tag path. supports_structural_tag() + # is True so required/named tool_choice routes through FunctionCallParser + # instead of JsonArrayParser. + + def test_supports_structural_tag(self): + self.assertTrue(self.detector.supports_structural_tag()) + + def test_get_model_structural_tag(self): + import xgrammar as xgr + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + tool_choice_name = ToolChoiceFuncName(name="get_current_weather") + tool_choice = ToolChoice(function=tool_choice_name) + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + +class TestGptOssDetector(unittest.TestCase): + def setUp(self): + self.tools = [ + Tool( + type="function", + function=Function( + name="search", + description="Searches for information.", + parameters={ + "type": "object", + "properties": { + "query": {"type": "string"}, + "topn": {"type": "integer"}, + }, + "required": ["query"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather information for a city.", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string"}, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["city"], + }, + ), + ), + ] + self.detector = GptOssDetector() + + def test_get_model_structural_tag(self): + import xgrammar as xgr + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice="required" + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + tool_choice_name = ToolChoiceFuncName(name="search") + tool_choice = ToolChoice(function=tool_choice_name) + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=True, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + + structural_tag = self.detector.get_structural_tag( + self.tools, thinking_mode=False, tool_choice=tool_choice + ) + self.assertIsInstance(structural_tag, xgr.StructuralTag) + grammar = xgr.Grammar.from_structural_tag(structural_tag) + self.assertIsInstance(grammar, xgr.Grammar) + class TestGlm4MoeDetector(unittest.TestCase): def setUp(self): @@ -3978,6 +4599,23 @@ class TestGetStructureConstraint(unittest.TestCase): self.assertEqual(result[0], "structural_tag") self.assertFalse(result[1].at_least_one) + def test_kimi_routes_through_legacy_with_section_markers(self): + """xgrammar 0.2.0's get_kimi_structural_tag(tool_choice='auto') emits + a bare <|tool_call_begin|>...<|tool_call_end|> grammar without the + section wrapper Kimi's chat template uses, so the parser would drop + any generated tool calls. KimiK2Detector therefore stays on the + legacy path; pin that here so a future tweak doesn't silently + re-route Kimi through the broken builtin.""" + from sglang.srt.entrypoints.openai.protocol import ( + LegacyStructuralTagResponseFormat, + ) + + parser = self._make_parser("kimi_k2", strict=True) + result = parser.get_structure_constraint("auto") + self.assertIsInstance(result[1], LegacyStructuralTagResponseFormat) + self.assertIn("<|tool_calls_section_begin|>", result[1].structures[0].begin) + self.assertIn("<|tool_calls_section_end|>", result[1].structures[0].end) + def test_kimi_auto_no_strict_returns_none(self): """auto without strict should not constrain.""" parser = self._make_parser("kimi_k2", strict=False) @@ -4015,28 +4653,38 @@ class TestGetStructureConstraint(unittest.TestCase): parser = self._make_parser("kimi_k2", strict=True) result = parser.get_structure_constraint("required") tag = result[1] - structures = tag.structures - self.assertTrue(len(structures) > 0) - self.assertIn("<|tool_calls_section_begin|>", structures[0].begin) - self.assertIn("<|tool_call_end|>", structures[0].end) + self.assertTrue(len(tag.structures) > 0) + self.assertIn("<|tool_calls_section_begin|>", tag.structures[0].begin) + self.assertIn("<|tool_call_end|>", tag.structures[0].end) def test_kimi_required_no_strict_uses_empty_schema(self): """Without strict, structural_tag should use empty schema per OpenAI protocol: strict=False means no parameter schema enforcement.""" parser = self._make_parser("kimi_k2", strict=False) result = parser.get_structure_constraint("required") - tag = result[1] - self.assertEqual(tag.structures[0].schema_, {}) + self.assertEqual(result[1].structures[0].schema_, {}) def test_kimi_required_strict_uses_tool_schema(self): """With strict, structural_tag should include the tool's parameter schema.""" parser = self._make_parser("kimi_k2", strict=True) result = parser.get_structure_constraint("required") - tag = result[1] - schema = tag.structures[0].schema_ + schema = result[1].structures[0].schema_ self.assertIn("properties", schema) self.assertIn("city", schema["properties"]) + # --- reasoning-prefix ownership --- + + def test_default_thinking_mode_is_false(self): + """Default must be False so callers don't silently get a reasoning + prefix added to their grammar (only relevant for detectors routed + through the xgrammar builtin).""" + import inspect + + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + sig = inspect.signature(FunctionCallParser.get_structure_constraint) + self.assertIs(sig.parameters["thinking_mode"].default, False) + class TestQwen25Detector(unittest.TestCase): """Test Qwen25Detector streaming and non-streaming multi-tool-call parsing.""" diff --git a/test/registered/unit/spec/test_spec_utils_traverse_tree.py b/test/registered/unit/spec/test_spec_utils_traverse_tree.py new file mode 100644 index 000000000..d580e024c --- /dev/null +++ b/test/registered/unit/spec/test_spec_utils_traverse_tree.py @@ -0,0 +1,71 @@ +"""Regression test for spec_utils.traverse_tree calling xgrammar with tensors. + +xgrammar 0.2.0 tightened its FFI binding and rejects 0-d tensors where Python +ints are expected. The dfs in traverse_tree recurses with `retrieve_next_token[curr]` +and reads `draft_tokens[curr]`, both of which return 0-d tensors and must be +explicitly cast before being handed to the grammar matcher. +""" + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.srt.speculative.spec_utils import traverse_tree +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=4, suite="stage-a-test-cpu") + + +class TestTraverseTreePassesIntsToGrammar(unittest.TestCase): + def _record_grammar(self): + """A grammar mock that records every call argument and rejects torch tensors.""" + grammar = MagicMock() + grammar.is_terminated.return_value = False + accept_calls = [] + fill_calls = [] + + def record_accept(token): + if isinstance(token, torch.Tensor): + raise TypeError(f"accept_token got torch.Tensor: {token!r}") + accept_calls.append(token) + + def record_fill(bitmask, idx): + if isinstance(idx, torch.Tensor): + raise TypeError(f"fill_vocab_mask got torch.Tensor idx: {idx!r}") + fill_calls.append(idx) + + grammar.accept_token.side_effect = record_accept + grammar.fill_vocab_mask.side_effect = record_fill + grammar.rollback.return_value = None + return grammar, accept_calls, fill_calls + + def test_branching_tree_passes_ints(self): + # Binary tree exercises both child recursion and sibling recursion: + # 0 ─┬─ 1 + # └─ 2 ─── 3 + retrieve_next_token = torch.tensor([1, -1, 3, -1], dtype=torch.int32) + retrieve_next_sibling = torch.tensor([-1, 2, -1, -1], dtype=torch.int32) + draft_tokens = torch.tensor([100, 11, 22, 33], dtype=torch.int64) + # all bits set: every draft token passes the parent's bitmask check + bitmask = torch.full((4, 4), -1, dtype=torch.int32) + + grammar, accept_calls, fill_calls = self._record_grammar() + traverse_tree( + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + grammar, + bitmask, + ) + + self.assertEqual(set(accept_calls), {11, 22, 33}) + self.assertEqual(set(fill_calls), {0, 1, 2, 3}) + for token in accept_calls: + self.assertIsInstance(token, int) + for idx in fill_calls: + self.assertIsInstance(idx, int) + + +if __name__ == "__main__": + unittest.main()