diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index 52b890f1cf11..2e234659c34c 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -210,6 +210,16 @@ These options are set within the YAML file passed to `trtllm-serve` via the `--c * **Description:** Required to load the Kimi K3 configuration and tokenizer code shipped with the checkpoint. +### Kimi-Specific API Behavior and Environment Variables + +When the served model is Kimi K3, `trtllm-serve` applies Kimi/Moonshot API semantics on `/v1/chat/completions` (all of these are exercised by Moonshot's [Kimi Vendor Verifier](https://www.kimi.com/blog/kimi-vendor-verifier.html)): + +* **Request extensions:** the `thinking` object (`{"type": "enabled"|"disabled", "keep": "all", "effort": "low"|"high"|"max"}`), `reasoning_effort` values `"low"`, `"high"`, `"max"`, and `"none"` (an explicit `thinking` object takes precedence), `tool_choice: "required"`, message-level (dynamic) tools declared on system messages, and `response_format` `json_object`/`json_schema` (the `json_schema` wrapper must carry a non-empty `name` and a `schema` object). These map onto the checkpoint chat template's native control messages; explicit `chat_template_kwargs` always win. +* **Streaming usage:** `usage` is reported in the final streaming chunk even when the client does not send `stream_options` (Kimi API parity). +* **Prompt-token accounting:** reported `usage.prompt_tokens` excludes the trailing 3-token generation channel opener, matching Kimi's reference accounting; the model still consumes the full rendered prompt. +* **`TRTLLM_KIMI_PARAM_POLICY`** (default `0`, off): when set to `1`, enforces Kimi's immutable sampling parameters — `top_p` pinned to 0.95 (unset or the OpenAI default `1.0` are coerced to 0.95; other values are rejected with HTTP 400), `presence_penalty`/`frequency_penalty` 0, `n` 1, and `temperature` bounded to [0, 1]. Off by default so existing deployments keep accepting the requests they accept today; a Kimi-Vendor-Verifier certification run must set it to `1` (the KVV params suite requires the rejections). +* **`TRTLLM_KIMI_K3_STRICT_TOOL_GRAMMAR`** (default `0`): opt-in constrained decoding for tools with `strict: true` (requires `guided_decoding_backend: xgrammar`). Disabled by default pending the investigation of a device-side assert observed under sustained concurrent guided load; strict tools otherwise fall back to warn-and-continue. + ## Testing API Endpoint The server (the OpenAI-compatible REST endpoint) runs on the rank-0 node, listening on port `8000`. Send requests to that node's hostname or IP; `localhost` only works from the rank-0 node itself. diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index 9b4ce83df05a..3834cbdb2c56 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -443,6 +443,9 @@ class KimiK3InputProcessor(KimiK25InputProcessor): "image": "<|kimi_image_placeholder|>", }, placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + # K3's reference renderer concatenates content parts with no + # separator; the default "\n" join skews prompt-token parity. + placeholders_separator="", ), ) class KimiK3ForConditionalGeneration(KimiK25ForConditionalGeneration): diff --git a/tensorrt_llm/executor/postproc_worker.py b/tensorrt_llm/executor/postproc_worker.py index 1981775d5299..184b9f924c8e 100644 --- a/tensorrt_llm/executor/postproc_worker.py +++ b/tensorrt_llm/executor/postproc_worker.py @@ -32,6 +32,12 @@ class PostprocArgs: first_iteration: bool = True num_prompt_tokens: Optional[int] = None + # Subtracted from num_prompt_tokens when reporting usage. Lets servers + # exclude generation-prompt stub tokens the vendor accounting treats as + # pending output (e.g. Kimi K3's 3-token channel opener) without changing + # what the model sees. num_prompt_tokens itself is overwritten by the + # executor on the postproc-worker path, so the offset must be separate. + num_prompt_tokens_offset: int = 0 tokenizer: Optional[TransformersTokenizer] = None ctx_usage: Optional[Any] = None diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 8a6d406604be..eac607673693 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -351,11 +351,15 @@ class ConversationMessage(TypedDict, total=False): This is used by `interleave_mm_placeholders` to insert multimodal placeholders at the correct positions, and to reconstruct the OpenAI-style content list for templates that handle media natively. + tools: Message-level (dynamic) tool declarations carried on system messages. Only + populated for models whose python-renderer chat template consumes them (kimi_k3); + absent for other models. """ role: str content: str media: List[MultimodalData] content_parts: List[Union[str, dict]] + tools: List[Dict[str, Any]] class MultimodalDataTracker: diff --git a/tensorrt_llm/serve/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index d60d93f2ce45..29dc1c69f7a1 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -247,7 +247,9 @@ def parse_chat_message_content_parts( def parse_chat_message_content( message: ChatCompletionMessageParam, - mm_data_tracker: MultimodalDataTracker) -> ConversationMessage: + mm_data_tracker: MultimodalDataTracker, + lenient_tool_call_arguments: bool = False, + keep_message_tools: bool = False) -> ConversationMessage: """Parse the content of a chat message.""" role = message["role"] content = message.get("content") @@ -265,9 +267,17 @@ def parse_chat_message_content( mm_data_tracker, ) if role == "assistant": - result.update(_parse_assistant_message_content(message)) + result.update( + _parse_assistant_message_content(message, + lenient_tool_call_arguments)) elif role == "tool": result.update(_parse_tool_message_content(message)) + elif keep_message_tools and role == "system" and message.get("tools"): + # Message-level (dynamic) tool declarations: python-renderer chat + # templates (kimi_k3) render these as an in-conversation tool + # declare block at this message's position. Other models keep the + # pre-existing behavior of silently ignoring the key. + result["tools"] = message["tools"] return result @@ -329,7 +339,10 @@ def _validate_fallback_tool_calls( return [tc.model_dump(exclude_unset=True) for tc in validated] -def _normalize_tool_call_arguments(index: int, item: Any) -> dict[str, Any]: +def _normalize_tool_call_arguments(index: int, + item: Any, + lenient_json: bool = False + ) -> dict[str, Any]: """Normalize `function.arguments` to the internal dict form.""" item = dict(item) item["function"] = dict(item["function"]) @@ -341,6 +354,11 @@ def _normalize_tool_call_arguments(index: int, item: Any) -> dict[str, Any]: try: arguments = json.loads(arguments) except json.JSONDecodeError as e: + if lenient_json: + # Keep the raw string: python-renderer templates (kimi_k3) + # normalize unparsable arguments themselves and render them + # verbatim as a JSON block, matching the reference tokenizer. + return item raise ValueError( f"tool_calls[{index}].function.arguments must be valid JSON." ) from e @@ -357,7 +375,9 @@ def _normalize_tool_call_arguments(index: int, item: Any) -> dict[str, Any]: return item -def _parse_fallback_tool_calls(tool_calls: list[Any]) -> list[dict[str, Any]]: +def _parse_fallback_tool_calls( + tool_calls: list[Any], + lenient_json: bool = False) -> list[dict[str, Any]]: """Parse raw tool-call lists accepted only by the tau2-bench fallback path. `openai_server.py` first attempts strict OpenAI request validation. Some tau2-bench requests @@ -369,13 +389,15 @@ def _parse_fallback_tool_calls(tool_calls: list[Any]) -> list[dict[str, Any]]: """ tool_calls = _validate_fallback_tool_calls(tool_calls) return [ - _normalize_tool_call_arguments(index, item) + _normalize_tool_call_arguments(index, item, lenient_json) for index, item in enumerate(tool_calls) ] # Adapted from: https://github.com/vllm-project/vllm/blob/4574d48bab9c4e38b7c0a830eeefc8f0980e8c58/vllm/entrypoints/chat_utils.py#L1406 -def _parse_assistant_message_content(message: Dict[str, Any]) -> Dict[str, Any]: +def _parse_assistant_message_content(message: Dict[str, Any], + lenient_json: bool = False + ) -> Dict[str, Any]: result = {} # Include reasoning if present for interleaved thinking. reasoning_content = message.get("reasoning") @@ -387,13 +409,14 @@ def _parse_assistant_message_content(message: Dict[str, Any]) -> Dict[str, Any]: tool_calls = message.get("tool_calls") if tool_calls is not None: if isinstance(tool_calls, list): - result["tool_calls"] = _parse_fallback_tool_calls(tool_calls) + result["tool_calls"] = _parse_fallback_tool_calls( + tool_calls, lenient_json) else: # The strict parse path delivers tool_calls as a single-use Pydantic `ValidatorIterator` # of already-validated OpenAI tool calls, so only materialize and normalize arguments. tool_calls = list(tool_calls) result["tool_calls"] = [ - _normalize_tool_call_arguments(index, item) + _normalize_tool_call_arguments(index, item, lenient_json) for index, item in enumerate(tool_calls) ] @@ -478,7 +501,13 @@ def parse_chat_messages_coroutines( content_format = ContentFormat.STRING for msg in messages: - parsed_msg = parse_chat_message_content(msg, mm_data_tracker) + # kimi_k3's reference renderer keeps unparsable tool-call argument + # strings verbatim; other templates expect the strict dict contract. + parsed_msg = parse_chat_message_content( + msg, + mm_data_tracker, + lenient_tool_call_arguments=(model_type == "kimi_k3"), + keep_message_tools=(model_type == "kimi_k3")) conversation.append(parsed_msg) # Track placeholders added for this message only. diff --git a/tensorrt_llm/serve/harmony_adapter.py b/tensorrt_llm/serve/harmony_adapter.py index 4b507108575b..392d0e1b8f1a 100644 --- a/tensorrt_llm/serve/harmony_adapter.py +++ b/tensorrt_llm/serve/harmony_adapter.py @@ -1977,12 +1977,17 @@ def _create_usage_info(num_prompt_tokens, def maybe_transform_reasoning_effort( - reasoning_effort: ReasoningEffort | Literal["low", "medium", "high"] | None + reasoning_effort: ReasoningEffort + | Literal["low", "medium", "high", "max", "none"] | None ) -> ReasoningEffort | None: str_to_effort = { "low": ReasoningEffort.LOW, "medium": ReasoningEffort.MEDIUM, - "high": ReasoningEffort.HIGH + "high": ReasoningEffort.HIGH, + # Kimi-style efforts accepted by the shared request schema; map to + # the nearest harmony level ("none" means no explicit effort). + "max": ReasoningEffort.HIGH, + "none": None, } if reasoning_effort and not isinstance(reasoning_effort, ReasoningEffort): return str_to_effort[reasoning_effort] diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 13f8c4bec5b3..878971c30080 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -409,6 +409,7 @@ class EmbeddingResponse(OpenAIBaseModel): def _response_format_to_guided_decoding_params( response_format: Optional[ResponseFormat], reasoning_parser: Optional[str] = None, + chat_template_kwargs: Optional[Dict[str, Any]] = None, ) -> Optional[GuidedDecodingParams]: if response_format is None: guided_decoding_params = None @@ -490,6 +491,32 @@ def _response_format_to_guided_decoding_params( "stop_after_first": True, } + elif reasoning_parser == "kimi_k3": + # K3 XTML: the generation prompt already ends inside the channel the + # model starts in. In thinking mode (the default) the response channel + # opens mid-generation, so trigger the user constraint on it + # (mirrors the gpt_oss final-channel handling). In non-thinking mode + # the prompt ends inside <|open|>response<|sep|>, the trigger would + # never be generated, and the raw grammar applies from the first + # generated token instead. + thinking = (chat_template_kwargs or {}).get("thinking", + True) is not False + if not thinking: + return guided_decoding_params + stag_format = { + "type": + "triggered_tags", + "triggers": ["<|open|>response<|sep|>"], + "tags": [ + { + "begin": "<|open|>response<|sep|>", + "content": content, + "end": "<|close|>response<|sep|>", + }, + ], + "stop_after_first": + True, + } else: # Force thinking and then trigger user constraint parser = ReasoningParserFactory.create_reasoning_parser( @@ -793,7 +820,24 @@ class ReasoningAssistantMessage(ChatCompletionAssistantMessageParam): reasoning_content: Optional[str] -ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam, +class DynamicToolsSystemMessageParam(TypedDict, total=False): + """System message carrying message-level (dynamic) tool declarations. + + Kimi-style templates render such messages as an in-conversation tool + declare block. Must come first in `ChatCompletionMessageParam`: the + stock OpenAI system-message TypedDict otherwise wins smart-union scoring + and silently drops the `tools` key. + """ + __pydantic_config__ = ConfigDict(extra="allow") # type: ignore + + role: Required[Literal["system"]] + tools: Required[List[dict]] + content: Union[str, List[ChatCompletionContentPartParam], None] + name: str + + +ChatCompletionMessageParam = Union[DynamicToolsSystemMessageParam, + OpenAIChatCompletionMessageParam, CustomChatCompletionMessageParam, ReasoningAssistantMessage] @@ -879,6 +923,19 @@ class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): type: Literal["function"] = "function" +class ChatCompletionThinkingParam(OpenAIBaseModel): + """Kimi/Moonshot `thinking` extension controlling reasoning output. + + `keep` is fixed to `"all"` when thinking is enabled and ignored when + disabled; `effort` is only meaningful when enabled. An explicit + `effort` wins over a request-level `reasoning_effort`, which applies + only when `effort` is absent (and never for a disabled request). + """ + type: Literal["enabled", "disabled"] = "enabled" + keep: Optional[Literal["all"]] = None + effort: Optional[Literal["low", "high", "max"]] = None + + class ChatCompletionRequest(OpenAIBaseModel): # Ordered by official OpenAI API documentation # https://platform.openai.com/docs/api-reference/chat/create @@ -906,17 +963,22 @@ class ChatCompletionRequest(OpenAIBaseModel): temperature: Optional[float] = None top_p: Optional[float] = None tools: Optional[List[ChatCompletionToolsParam]] = None - tool_choice: Optional[Union[Literal["none", "auto"], + tool_choice: Optional[Union[Literal["none", "auto", "required"], ChatCompletionNamedToolChoiceParam]] = "none" user: Optional[str] = None reasoning_effort: Optional[ReasoningEffort | Literal[ - "low", "medium", "high"]] = Field( + "low", "medium", "high", "max", "none"]] = Field( default=ReasoningEffort.LOW, description=( "The level of reasoning effort to use. Controls how much " "reasoning is shown in the model's response. Options: " - "'low', 'medium', 'high'."), + "'low', 'medium', 'high' (harmony/gpt-oss), plus 'max' and " + "'none' for models with Kimi-style thinking control."), ) + # Kimi/Moonshot extension: structured control of reasoning output. The + # serving layer maps it into the chat-template kwargs for models whose + # template understands it (e.g. kimi_k3); other models ignore it. + thinking: Optional[ChatCompletionThinkingParam] = None thinking_token_budget: Optional[int] = None prompt_ignore_length: Optional[int] = 0 @@ -1086,7 +1148,9 @@ def to_sampling_params(self, spaces_between_special_tokens=self.spaces_between_special_tokens, truncate_prompt_tokens=self.truncate_prompt_tokens, guided_decoding=_response_format_to_guided_decoding_params( - self.response_format, reasoning_parser=reasoning_parser), + self.response_format, + reasoning_parser=reasoning_parser, + chat_template_kwargs=self.chat_template_kwargs), thinking_token_budget=self.thinking_token_budget, # logits_bias @@ -1112,14 +1176,47 @@ def validate_stream_options(cls, values): @model_validator(mode="before") @classmethod def check_tool_choice(cls, data): - if "tool_choice" not in data and data.get("tools"): + if not isinstance(data, dict): + return data + has_dynamic_tools = any( + isinstance(msg, dict) and msg.get("role") == "system" + and msg.get("tools") for msg in data.get("messages") or []) + if "tool_choice" not in data and (data.get("tools") + or has_dynamic_tools): data["tool_choice"] = "auto" - if "tool_choice" in data and data["tool_choice"] != "none": - if "tools" not in data or data["tools"] is None: + # "none" and "auto" are meaningful without tools. "required" needs a + # non-empty tool set — request-level or message-level (dynamic) + # tools both count. A named function must be a request-level tool. + if "tool_choice" in data and data["tool_choice"] not in ("none", + "auto"): + satisfied = data.get("tools") or (data["tool_choice"] == "required" + and has_dynamic_tools) + if not satisfied: raise ValueError( "When using `tool_choice`, `tools` must be set.") return data + @model_validator(mode="before") + @classmethod + def check_message_tools_role(cls, data): + """Message-level tool declarations ride on system messages only. + + Union validation strips unknown keys from non-system messages, so + this raw-payload validator is the only layer that can reject the + misuse loudly instead of silently dropping the client's tools. The + rest of the dynamic-tools contract is model-specific and validated + in the serving layer (kimi_k3 only). + """ + if not isinstance(data, dict): + return data + for message in data.get("messages") or []: + if (isinstance(message, dict) and message.get("tools") is not None + and message.get("role") != "system"): + raise ValueError( + "Message-level `tools` are only allowed on system " + "messages.") + return data + @model_validator(mode="before") @classmethod def check_logprobs(cls, data): diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 8e0622a6725f..723804b4fd30 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -72,15 +72,16 @@ QueueFullError) from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_protocol import ( - ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, - ChatCompletionResponse, ChatCompletionResponseChoice, ChatMessage, - CompletionRequest, CompletionResponse, CompletionResponseChoice, - EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, - EmbeddingUsageInfo, ErrorResponse, ImageEditRequest, ImageGenerationRequest, - ImageGenerationResponse, ImageObject, MemoryUpdateRequest, ModelCard, - ModelList, PromptTokensDetails, ResponseFormat, ResponsesRequest, - ResponsesResponse, TokenizeRequest, TokenizeResponse, UpdateWeightsRequest, - UsageInfo, ensure_request_chat_template_allowed, to_llm_conversation_params, + ChatCompletionMessageParam, ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, + ChatCompletionToolsParam, ChatMessage, CompletionRequest, + CompletionResponse, CompletionResponseChoice, EmbeddingRequest, + EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, + ImageEditRequest, ImageGenerationRequest, ImageGenerationResponse, + ImageObject, MemoryUpdateRequest, ModelCard, ModelList, PromptTokensDetails, + ResponseFormat, ResponsesRequest, ResponsesResponse, StreamOptions, + TokenizeRequest, TokenizeResponse, UpdateWeightsRequest, UsageInfo, + ensure_request_chat_template_allowed, to_llm_conversation_params, to_llm_disaggregated_params) from tensorrt_llm.serve.openai_video_routes import _VideoRoutesMixin from tensorrt_llm.serve.perf_metrics import (PerfMetricsJsonlWriter, @@ -206,6 +207,183 @@ def _warn_unresolvable_thinking_once(reasoning_parser: str) -> None: "build that relays 'resolved_thinking'.") +def _enforce_kimi_param_policy(request: ChatCompletionRequest) -> None: + """Enforce Kimi's immutable sampling-parameter policy (KVV params suite). + + Kimi's API pins top_p, the penalties, and n, and bounds temperature to + [0, 1]; out-of-policy values must fail fast with HTTP 400 rather than + generate. top_p unset or the OpenAI-default 1.0 is coerced to the pinned + 0.95 instead of rejected. Off by default so existing K3 deployments keep + accepting the requests they accept today (review feedback); a Kimi + Vendor Verifier certification run must opt in with + TRTLLM_KIMI_PARAM_POLICY=1. + """ + if os.getenv("TRTLLM_KIMI_PARAM_POLICY", "0") != "1": + return + if request.top_p is None or request.top_p == 1.0: + # Kimi pins top_p at 0.95. None would fall back to 1.0 in + # to_sampling_params; an explicit 1.0 is the OpenAI SDK default many + # clients send unconditionally — coerce both to the pinned value + # rather than rejecting (review feedback). + request.top_p = 0.95 + if request.temperature is not None and not (0.0 <= request.temperature <= + 1.0): + raise ValueError("temperature must be within [0, 1] for this model; " + f"got {request.temperature}.") + if request.top_p is not None and request.top_p != 0.95: + raise ValueError( + f"top_p is fixed at 0.95 for this model; got {request.top_p}.") + if request.presence_penalty: + raise ValueError("presence_penalty is fixed at 0 for this model; " + f"got {request.presence_penalty}.") + if request.frequency_penalty: + raise ValueError("frequency_penalty is fixed at 0 for this model; " + f"got {request.frequency_penalty}.") + if request.n != 1: + raise ValueError(f"n is fixed at 1 for this model; got {request.n}.") + + +# Valid function-tool name: no leading digit, word chars/dash only, at most +# 256 chars (Kimi Vendor Verifier contract for message-level tools). +_DYNAMIC_TOOL_NAME_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_-]{0,255}\Z") + + +def _dynamic_tool_dicts( + messages: Optional[List[ChatCompletionMessageParam]]) -> list[dict]: + """Collect message-level (dynamic) tool declarations from system messages.""" + tools: list[dict] = [] + for msg in messages or []: + if isinstance( + msg, dict) and msg.get("role") == "system" and msg.get("tools"): + tools.extend(msg["tools"]) + return tools + + +def _validate_kimi_dynamic_tools(request: ChatCompletionRequest) -> None: + """Validate message-level (dynamic) tool declarations for kimi_k3. + + Kimi-style dynamic tools ride on system messages. Enforce the contract + checked by the Kimi Vendor Verifier: system-only carrier, empty content, + function-typed tools with valid unique names (unique also against + request-level tools). Only called for kimi_k3 deployments; other models + keep ignoring the key as before. + """ + seen_names = set() + for tool in request.tools or []: + seen_names.add(tool.function.name) + for message in request.messages or []: + # A null tools key is treated as absent (some SDKs serialize + # optional fields as null); only declared tools are validated. + if not isinstance(message, dict) or message.get("tools") is None: + continue + if message.get("role") != "system": + raise ValueError( + "Message-level `tools` are only allowed on system messages.") + if message.get("content"): + raise ValueError( + "A system message carrying `tools` must have empty content.") + message_tools = message["tools"] + if not isinstance(message_tools, list): + raise ValueError("Message-level `tools` must be an array.") + for tool in message_tools: + if not isinstance(tool, dict): + raise ValueError("Each message-level tool must be an object.") + if tool.get("type") != "function": + raise ValueError(f"Unsupported message-level tool type: " + f"{tool.get('type')!r}.") + function = tool.get("function") + if not isinstance(function, dict): + raise ValueError( + "Message-level tools must carry a `function` object.") + name = function.get("name") + if not isinstance(name, + str) or not _DYNAMIC_TOOL_NAME_RE.match(name): + raise ValueError(f"Invalid message-level tool name: {name!r}.") + if name in seen_names: + raise ValueError(f"Duplicate tool name: {name!r}.") + seen_names.add(name) + + +def _apply_kimi_chat_extensions(request: ChatCompletionRequest, + model_type: Optional[str]) -> None: + """Apply Kimi/Moonshot API semantics to a chat request for kimi_k3. + + The kimi_k3 checkpoint template natively renders control messages for + thinking effort, tool_choice, and response_format, but only reads them + from chat-template kwargs. Derive those kwargs from the request-level + fields so the OpenAI-style API surface drives the template; explicit + client-supplied `chat_template_kwargs` win over derived values. The + merged kwargs also steer the kimi_k3 reasoning parser's initial channel, + the guided-decoding structural tag, and the thinking-budget logits + processor downstream. + + Kimi's API also reports usage in the final streaming chunk without the + client opting in, so default `stream_options` for streaming requests. + """ + if model_type != "kimi_k3": + return + _validate_kimi_dynamic_tools(request) + _enforce_kimi_param_policy(request) + if request.stream and request.stream_options is None: + # StreamOptions defaults: include_usage=True, continuous off. + request.stream_options = StreamOptions() + derived: dict[str, Any] = {} + if request.thinking is not None: + enabled = request.thinking.type != "disabled" + derived["thinking"] = enabled + if enabled and request.thinking.effort is not None: + derived["thinking_effort"] = request.thinking.effort + if ("reasoning_effort" in request.model_fields_set + and request.reasoning_effort is not None + and "thinking_effort" not in derived and + (request.thinking is None or request.thinking.type != "disabled")): + # Kimi semantics: an explicit thinking.effort wins, and an explicit + # thinking object also wins the on/off axis — reasoning_effort only + # supplies the effort when thinking.effort is absent, and + # reasoning_effort="none" only disables thinking when no thinking + # object was sent. No effort is ever derived for an explicitly + # disabled request. (KVV test_reasoning_effort_ignored_when_effort_ + # present / test_reasoning_effort_effective_when_effort_absent.) + effort = getattr(request.reasoning_effort, "value", + request.reasoning_effort).lower() + if effort == "none": + if request.thinking is None: + derived["thinking"] = False + elif effort in ("low", "high", "max"): + derived["thinking_effort"] = effort + # Other efforts (e.g. harmony's "medium") have no K3 equivalent; + # leave the template default. + if ("tool_choice" in request.model_fields_set + and (request.tools or _dynamic_tool_dicts(request.messages)) + and request.tool_choice in ("required", "none")): + derived["tool_choice"] = request.tool_choice + response_format = request.response_format + if response_format is not None and response_format.type in ("json_object", + "json_schema"): + derived["response_format"] = response_format.type + if response_format.type == "json_schema": + # Kimi requires the OpenAI wrapper shape: {name, schema[, strict]}. + json_schema = response_format.json_schema + if not isinstance(json_schema, dict) or not isinstance( + json_schema.get("name"), str) or not json_schema["name"]: + raise ValueError( + "response_format.json_schema requires a non-empty " + "`name` string.") + if not isinstance(json_schema.get("schema"), dict): + raise ValueError( + "response_format.json_schema requires a `schema` object.") + if "strict" in json_schema and not isinstance( + json_schema["strict"], bool): + raise ValueError( + "response_format.json_schema.strict must be a boolean.") + derived["response_schema"] = json_schema["schema"] + if derived: + request.chat_template_kwargs = { + **derived, + **(request.chat_template_kwargs or {}), + } + + def _configure_parser_special_token_decoding( sampling_params: SamplingParams, reasoning_parser_name: Optional[str], tool_parser_name: Optional[str], has_tools: bool) -> None: @@ -267,6 +445,14 @@ def _build_tool_strict_guided_decoding_params(tools, tool_parser_name): return None parser = tool_parser_cls() + # Parsers whose wire format can't be expressed through structure_info + # triples (kimi_k3 XTML) provide a complete format themselves. + custom_format = parser.build_strict_structural_tag_format(tools) + if custom_format is not None: + resp_format = ResponseFormat(type="structural_tag", + format=custom_format) + return GuidedDecodingParams(structural_tag=resp_format.model_dump_json( + by_alias=True, exclude_none=True)) if not parser.supports_structural_tag(): logger.warning( "Tool parser '%s' does not support structural tags, " @@ -1673,9 +1859,23 @@ async def chat_stream_generator( try: ensure_request_chat_template_allowed( request, self.allow_request_chat_template) + model_type = resolve_top_level_model_type(self.model_config) + is_kimi_k3 = model_type == "kimi_k3" + _apply_kimi_chat_extensions(request, model_type) + if request.tool_choice == "required" and not is_kimi_k3: + # Schema-accepting "required" everywhere but enforcing it only + # for kimi_k3 would silently degrade to "auto" elsewhere; + # reject loudly for models that cannot honor it. + raise ValueError( + "tool_choice='required' is not supported for this model.") conversation: List[ConversationMessage] = [] + # exclude_none for kimi_k3: pydantic-injected null defaults + # (strict, description, parameters) would otherwise leak into the + # rendered tool-declare JSON and skew prompt-token parity. Other + # models keep their historical rendering. tool_dicts = None if request.tools is None else [ - tool.model_dump() for tool in request.tools + tool.model_dump(exclude_none=is_kimi_k3) + for tool in request.tools ] # Pass the model vocabulary size so ``logit_bias`` can be # expanded into an embedding bias tensor in the sampler. @@ -1780,15 +1980,32 @@ async def chat_stream_generator( forced_tool_name = request.tool_choice.function.name reasoning_parser_name = self.generator.args.reasoning_parser + # Message-level (dynamic) tools are a Kimi API extension; only + # kimi_k3 templates render them, so other models keep ignoring + # the key entirely. + dynamic_tools = _dynamic_tool_dicts( + request.messages) if is_kimi_k3 else [] + dynamic_tool_params: List[ChatCompletionToolsParam] = [] + if dynamic_tools: + try: + dynamic_tool_params = [ + ChatCompletionToolsParam.model_validate(tool) + for tool in dynamic_tools + ] + except ValidationError as e: + raise ValueError( + f"Invalid message-level tool declaration: {e}") from e _configure_parser_special_token_decoding( sampling_params, reasoning_parser_name=reasoning_parser_name, tool_parser_name=self.tool_parser, - has_tools=bool(request.tools)) - if self.tool_parser and request.tools: + has_tools=bool(request.tools) or bool(dynamic_tools)) + all_tools = (request.tools or []) + dynamic_tool_params + if self.tool_parser and all_tools: # When strict=True on any tool, apply constrained decoding # via structural tags (only if response_format doesn't already - # set guided decoding). + # set guided decoding). Dynamic tools must be in the grammar + # too, or their calls would be masked out. if sampling_params.guided_decoding is None: if (forced_tool_name is not None and _parser_extracts_forced_tool_calls( @@ -1823,7 +2040,7 @@ async def chat_stream_generator( # ``triggered_tags`` semantics. Engages only when at # least one tool has ``strict=True``. strict_guided = _build_tool_strict_guided_decoding_params( - request.tools, self.tool_parser) + all_tools, self.tool_parser) if strict_guided is not None: sampling_params.guided_decoding = strict_guided elif forced_tool_name is not None: @@ -1834,6 +2051,25 @@ async def chat_stream_generator( err_type="BadRequestError", status_code=HTTPStatus.BAD_REQUEST) postproc_args = ChatPostprocArgs.from_request(request) + if (is_kimi_k3 and request.add_generation_prompt + and request.prompt_token_ids is None + and request.prompt_token_ids_b64 is None + and request.chat_template is None + and self.chat_template is None): + # Kimi's prompt-token accounting excludes the trailing 3-token + # generation channel opener (<|open|>think|response<|sep|>); + # the model still sees the full rendered prompt. b64-relayed + # token ids (decoded later) must behave like plain + # prompt_token_ids: no rendering here, so no stub to exclude. + # The offset presumes the checkpoint's native K3 renderer; an + # explicit request- or server-level chat template may end + # differently, so report unadjusted usage for those. + postproc_args.num_prompt_tokens_offset = 3 + if dynamic_tool_params: + # The tool parser must see dynamic tools to recognize their + # calls in the model output. + postproc_args.tools = (postproc_args.tools + or []) + dynamic_tool_params self._validate_internal_disagg_request(request, raw_request) disaggregated_params = to_llm_disaggregated_params( request.disaggregated_params) @@ -2395,6 +2631,11 @@ async def create_streaming_generator(promise: RequestOutput, "follow-up: harmony sub-task of TRTLLM-12758."), err_type="BadRequestError", status_code=HTTPStatus.BAD_REQUEST) + if request.tool_choice == "required": + # The harmony path treats unknown tool_choice values like + # "auto"; reject instead of silently degrading. + raise ValueError( + "tool_choice='required' is not supported for this model.") # Initialize HarmonyAdapter # NOTE: WAR for Disagg failure, may affect perf if no warmup if not self.harmony_adapter: diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 74b4515f8f11..27b30181f524 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -91,8 +91,10 @@ class ChatPostprocArgs(PostprocArgs): model: str num_choices: int = 1 tools: Optional[List[ChatCompletionToolsParam]] = None - tool_choice: Optional[Union[Literal["none"], - ChatCompletionNamedToolChoiceParam]] = "none" + # None means "not specified": only an explicit client "none" (always set + # by from_request) suppresses parsed tool calls in apply_tool_parser. + tool_choice: Optional[Union[Literal["none", "auto", "required"], + ChatCompletionNamedToolChoiceParam]] = None return_logprobs: bool = False top_logprobs: bool = False stream_options: Optional[StreamOptions] = None @@ -240,6 +242,11 @@ def apply_tool_parser(args: ChatPostprocArgs, result = StreamingParseResult( normal_text=result.normal_text + finish_result.normal_text, calls=result.calls + finish_result.calls) + if args.tool_choice == "none": + # tool_choice="none": still run the parser (including the finish + # flush above) so tool-call markup is stripped from content, but + # never surface tool calls. + return result.normal_text, [] normal_text, calls = result.normal_text, result.calls if result.calls: args.has_tool_call[output_index] = True @@ -347,7 +354,7 @@ def yield_first_chat(num_tokens: int, res: List[str] = [] finish_reason_sent = [False] * args.num_choices - prompt_tokens = args.num_prompt_tokens + prompt_tokens = args.num_prompt_tokens - args.num_prompt_tokens_offset ctx_usage = _ctx_usage_for_postproc(args, rsp.outputs) stream_response_id, stream_created = _ensure_stream_metadata( args, rsp, "chatcmpl") @@ -678,7 +685,7 @@ def chat_response_post_processor( full_message = args.last_message_content + choice.message.content choice.message.content = full_message - num_prompt_tokens = args.num_prompt_tokens + num_prompt_tokens = args.num_prompt_tokens - args.num_prompt_tokens_offset num_generated_tokens = sum(len(output.token_ids) for output in rsp.outputs) usage = UsageInfo( prompt_tokens=num_prompt_tokens, @@ -883,7 +890,7 @@ def completion_response_post_processor( class ChatCompletionPostprocArgs(PostprocArgs): model: str tools: Optional[List[ChatCompletionToolsParam]] - tool_choice: Optional[Union[Literal["none", "auto"], + tool_choice: Optional[Union[Literal["none", "auto", "required"], ChatCompletionNamedToolChoiceParam]] request_id: Optional[int] = None stream_options: Optional[StreamOptions] = None diff --git a/tensorrt_llm/serve/tool_parser/base_tool_parser.py b/tensorrt_llm/serve/tool_parser/base_tool_parser.py index 8cbea0be7893..2968169ababd 100644 --- a/tensorrt_llm/serve/tool_parser/base_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/base_tool_parser.py @@ -1,7 +1,7 @@ # Adapted from https://github.com/sgl-project/sglang/blob/083629c23564e1a64deaa052f1df5c5d914358d8/python/sglang/srt/function_call/base_format_detector.py import json from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from partial_json_parser.core.exceptions import MalformedJSON from partial_json_parser.core.options import Allow @@ -329,6 +329,17 @@ def supports_structural_tag(self) -> bool: """Return True if this detector supports structural tag format.""" return True + def build_strict_structural_tag_format( + self, tools: List[Tool]) -> Optional[Dict[str, Any]]: + """Build a complete structural-tag format for strict-tool decoding. + + Override on parsers whose wire format cannot be expressed through + the `structure_info` begin/end/trigger triples (e.g. kimi_k3's + XTML call tags). Returns the xgrammar structural-tag format dict, + or None to fall back to the `structure_info` path. + """ + return None + @abstractmethod def structure_info(self) -> _GetInfoFunc: """ diff --git a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py index c3998762b6b8..6291ccd6a735 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py @@ -28,6 +28,7 @@ """ import json +import os import re from typing import Any, Dict, List @@ -42,6 +43,10 @@ def _unescape_attr(value: str) -> str: return value.replace(""", '"').replace("&", "&") +def _escape_attr(value: str) -> str: + return value.replace("&", "&").replace('"', """) + + def _parse_attrs(header: str) -> Dict[str, str]: return {key: _unescape_attr(value) for key, value in re.findall(r'(\w+)="([^"]*)"', header)} @@ -96,8 +101,8 @@ def has_tool_call(self, text: str) -> bool: def supports_structural_tag(self) -> bool: # XTML argument bodies are tag-structured text, not JSON — the - # JSON-schema-driven structural-tag constrained decoding used for - # strict tools does not apply. + # generic begin/end/trigger structural-tag path does not apply. + # Strict tools are handled by build_strict_structural_tag_format. return False def structure_info(self) -> _GetInfoFunc: @@ -105,6 +110,98 @@ def structure_info(self) -> _GetInfoFunc: "kimi_k3 XTML tool calls do not support structural-tag constrained decoding" ) + def build_strict_structural_tag_format(self, tools: List[Tool]) -> Dict[str, Any] | None: + """Xgrammar structural-tag format enforcing well-formed K3 tool calls. + + Any generated tools section is constrained to calls of the declared + tools; a strict tool with a parameters schema additionally gets its + arguments constrained to that JSON Schema via the K3 json-block body + form (the per-argument XTML form has no xgrammar equivalent). + Non-strict tools keep free-form bodies. The outer triggered_tags + must keep `at_least_one`/`stop_after_first` False: True would + forbid the think/response text before the section and the message + close after it, deadlocking generation. + """ + if os.getenv("TRTLLM_KIMI_K3_STRICT_TOOL_GRAMMAR", "0") != "1": + # Experimental, opt-in: under concurrent guided load with + # production tool schemas, sampling tripped a device-side assert + # and hard-killed the deployment (KVV schema suite, job 3054205: + # TensorCompare.cu _assert_async in sampler.update_requests). + # Root-cause investigation pending; strict tools fall back to + # the warn-and-continue path meanwhile. + return None + if not tools: + return None + for tool in tools: + if "<" in tool.function.name: + # A literal '<' in an attribute value has no escaped form in + # the K3 wire format (the checkpoint renderer escapes only + # '&' and '"'), and the parser's attribute regex stops at + # '<' — a grammar-forced call with such a name would be + # dropped. Skip constrained decoding rather than teach the + # model a dialect the reference renderer never produces. + logger.warning( + f"Tool name {tool.function.name!r} contains '<'; " + "skipping the kimi_k3 strict-tool grammar for this request." + ) + return None + call_tags: List[Dict[str, Any]] = [] + for tool in tools: + begin = f'<|open|>call tool="{_escape_attr(tool.function.name)}"' + if tool.function.strict and tool.function.parameters: + call_tags.append( + { + "type": "tag", + "begin": begin, + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": ' index="[1-9][0-9]{0,2}"', + }, + { + "type": "const_string", + "value": '<|sep|><|open|>json type="object"<|sep|>', + }, + { + "type": "json_schema", + "json_schema": tool.function.parameters, + }, + ], + }, + "end": "<|close|>json<|sep|><|close|>call<|sep|>", + } + ) + else: + call_tags.append( + { + "type": "tag", + "begin": begin, + "content": {"type": "any_text"}, + "end": "<|close|>call<|sep|>", + } + ) + return { + "type": "triggered_tags", + "triggers": [self.bot_token], + "tags": [ + { + "type": "tag", + "begin": self.bot_token, + "content": { + "type": "tags_with_separator", + "separator": "", + "at_least_one": True, + "tags": call_tags, + }, + "end": self.eot_token, + } + ], + "at_least_one": False, + "stop_after_first": False, + } + @staticmethod def _coerce_value(value: str, value_type: str) -> Any: if value_type == "string": diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index f2998fc859b2..f50d7f05c3e0 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -76,6 +76,7 @@ l0_cpu: - unittest/llmapi/apps/test_chat_utils.py - unittest/llmapi/apps/test_harmony_channel_validation.py - unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting + - unittest/llmapi/apps/test_kimi_serve_extensions.py - unittest/llmapi/apps/test_reasoning_prompt_resolution.py - unittest/llmapi/apps/test_tool_parsers.py - unittest/llmapi/test_bench_async.py diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index f5846dc97b25..eae652bda660 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -612,7 +612,7 @@ models: required: false tool_choice: kind: openai - type: Optional[Union[Literal['none', 'auto'], ChatCompletionNamedToolChoiceParam]] + type: Optional[Union[Literal['none', 'auto', 'required'], ChatCompletionNamedToolChoiceParam]] default: none status: stable required: false @@ -624,10 +624,16 @@ models: required: false reasoning_effort: kind: openai - type: Optional[ReasoningEffort | Literal['low', 'medium', 'high']] + type: Optional[ReasoningEffort | Literal['low', 'medium', 'high', 'max', 'none']] default: "ReasoningEffort.LOW" status: stable required: false + thinking: + kind: extension + type: Optional[ChatCompletionThinkingParam] + default: null + status: prototype + required: false thinking_token_budget: kind: extension type: Optional[int] diff --git a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py new file mode 100644 index 000000000000..cfe784170985 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py @@ -0,0 +1,572 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU-only unit tests for the Kimi K3 serving extensions. + +Covers the Kimi Vendor Verifier (KVV) API contract implemented in +`tensorrt_llm/serve`: request validators, the Kimi extension-to-chat-template +mapping and its precedence rules, the immutable sampling-parameter policy, the +kimi_k3 `response_format` guided-decoding branch, and the (env-gated) +strict-tools structural-tag grammar builder. No GPU or checkpoint required. +""" + +import json + +import pytest +from pydantic import ValidationError + +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, + StreamOptions, + _response_format_to_guided_decoding_params, +) +from tensorrt_llm.serve.openai_server import ( + _apply_kimi_chat_extensions, + _dynamic_tool_dicts, + _enforce_kimi_param_policy, + _validate_kimi_dynamic_tools, +) +from tensorrt_llm.serve.tool_parser.kimi_k3_tool_parser import ( + KimiK3ToolParser, + _escape_attr, + _parse_attrs, + _unescape_attr, +) + +# The L0 CPU stage invokes registered files with `-m cpu_only`; without this +# marker every test is deselected and pytest exits with code 5. +pytestmark = pytest.mark.cpu_only + +WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather of a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + +USER_MSG = {"role": "user", "content": "what is the weather in beijing?"} + + +def make_request(**kwargs) -> ChatCompletionRequest: + kwargs.setdefault("model", "hf-kimi-k3") + kwargs.setdefault("messages", [USER_MSG]) + return ChatCompletionRequest(**kwargs) + + +def dynamic_system_msg(tools: list, **extra) -> dict: + return {"role": "system", "content": "", "tools": tools, **extra} + + +class TestToolChoiceValidation: + def test_required_with_tools_accepted(self) -> None: + req = make_request(tools=[WEATHER_TOOL], tool_choice="required") + assert req.tool_choice == "required" + + def test_required_without_tools_rejected(self) -> None: + with pytest.raises(ValidationError, match=r"tools.*must be set"): + make_request(tool_choice="required") + + def test_required_with_empty_tools_rejected(self) -> None: + with pytest.raises(ValidationError, match=r"tools.*must be set"): + make_request(tools=[], tool_choice="required") + + def test_required_with_dynamic_only_tools_accepted(self) -> None: + req = make_request( + messages=[dynamic_system_msg([WEATHER_TOOL]), USER_MSG], tool_choice="required" + ) + assert req.tool_choice == "required" + + def test_named_with_dynamic_only_tools_rejected(self) -> None: + with pytest.raises(ValidationError, match=r"tools.*must be set"): + make_request( + messages=[dynamic_system_msg([WEATHER_TOOL]), USER_MSG], + tool_choice={ + "type": "function", + "function": {"name": "get_weather"}, + }, + ) + + def test_auto_without_tools_accepted(self) -> None: + req = make_request(tool_choice="auto") + assert req.tool_choice == "auto" + + def test_tools_default_tool_choice_to_auto(self) -> None: + assert make_request(tools=[WEATHER_TOOL]).tool_choice == "auto" + + def test_dynamic_only_tools_default_tool_choice_to_auto(self) -> None: + req = make_request(messages=[dynamic_system_msg([WEATHER_TOOL]), USER_MSG]) + assert req.tool_choice == "auto" + + def test_no_tools_defaults_to_none(self) -> None: + assert make_request().tool_choice == "none" + + +class TestMessageToolsCarrierValidation: + """Carrier-role validation for message-level tools. + + The role restriction must reject at the raw-payload layer: union + validation strips unknown keys from non-system messages, so a + serving-side check would never see them. + """ + + def test_tools_in_user_message_rejected(self) -> None: + with pytest.raises(ValidationError, match="only allowed on system"): + make_request( + messages=[ + { + "role": "user", + "content": "hi", + "tools": [WEATHER_TOOL], + }, + USER_MSG, + ] + ) + + def test_tools_in_assistant_message_rejected(self) -> None: + with pytest.raises(ValidationError, match="only allowed on system"): + make_request( + messages=[ + USER_MSG, + { + "role": "assistant", + "content": "hello", + "tools": [WEATHER_TOOL], + }, + USER_MSG, + ] + ) + + def test_null_tools_key_ignored_everywhere(self) -> None: + req = make_request( + messages=[ + {"role": "user", "content": "hi", "tools": None}, + {"role": "system", "content": "be nice", "tools": None}, + USER_MSG, + ] + ) + assert _dynamic_tool_dicts(req.messages) == [] + + def test_system_tools_key_survives_validation(self) -> None: + req = make_request(messages=[dynamic_system_msg([WEATHER_TOOL]), USER_MSG]) + assert _dynamic_tool_dicts(req.messages) == [WEATHER_TOOL] + + def test_system_tools_with_content_survives_validation(self) -> None: + # Content correctness is enforced by the kimi-gated serving layer, + # but the key itself must not be silently stripped by the union. + req = make_request( + messages=[dynamic_system_msg([WEATHER_TOOL], content="not empty"), USER_MSG] + ) + assert _dynamic_tool_dicts(req.messages) == [WEATHER_TOOL] + + +class TestKimiDynamicToolsValidation: + def check(self, messages: list, **kwargs) -> None: + _validate_kimi_dynamic_tools(make_request(messages=messages, **kwargs)) + + def test_valid_dynamic_tool_passes(self) -> None: + self.check([dynamic_system_msg([WEATHER_TOOL]), USER_MSG]) + + def test_absent_content_passes(self) -> None: + self.check([{"role": "system", "tools": [WEATHER_TOOL]}, USER_MSG]) + + def test_strict_false_passes(self) -> None: + tool = json.loads(json.dumps(WEATHER_TOOL)) + tool["function"]["strict"] = False + self.check([dynamic_system_msg([tool]), USER_MSG]) + + def test_nonempty_content_rejected(self) -> None: + with pytest.raises(ValueError, match="empty content"): + self.check([dynamic_system_msg([WEATHER_TOOL], content="x"), USER_MSG]) + + def test_tools_not_array_rejected(self) -> None: + with pytest.raises(ValueError, match="must be an array"): + self.check( + [ + { + "role": "system", + "content": "", + "tools": {"type": "function"}, + }, + USER_MSG, + ] + ) + + def test_tool_item_not_object_rejected(self) -> None: + with pytest.raises(ValueError, match="must be an object"): + self.check([dynamic_system_msg([None]), USER_MSG]) + + def test_missing_type_rejected(self) -> None: + with pytest.raises(ValueError, match="Unsupported message-level"): + self.check([dynamic_system_msg([{"function": {"name": "x"}}]), USER_MSG]) + + def test_bogus_type_rejected(self) -> None: + with pytest.raises(ValueError, match="Unsupported message-level"): + self.check( + [dynamic_system_msg([{"type": "bogus", "function": {"name": "x"}}]), USER_MSG] + ) + + def test_missing_function_rejected(self) -> None: + with pytest.raises(ValueError, match="`function` object"): + self.check([dynamic_system_msg([{"type": "function"}]), USER_MSG]) + + @pytest.mark.parametrize( + "name", + [ + "1bad_name", + "bad@name", + "", + "a" * 257, + "get_weather\n", + ], + ) + def test_invalid_names_rejected(self, name: str) -> None: + with pytest.raises(ValueError, match="Invalid message-level"): + self.check( + [dynamic_system_msg([{"type": "function", "function": {"name": name}}]), USER_MSG] + ) + + @pytest.mark.parametrize("name", ["a" * 256, "Get_weather-2", "_x"]) + def test_valid_names_accepted(self, name: str) -> None: + self.check( + [dynamic_system_msg([{"type": "function", "function": {"name": name}}]), USER_MSG] + ) + + def test_duplicate_within_message_rejected(self) -> None: + with pytest.raises(ValueError, match="Duplicate tool name"): + self.check([dynamic_system_msg([WEATHER_TOOL, WEATHER_TOOL]), USER_MSG]) + + def test_duplicate_across_messages_rejected(self) -> None: + with pytest.raises(ValueError, match="Duplicate tool name"): + self.check( + [ + dynamic_system_msg([WEATHER_TOOL]), + dynamic_system_msg([WEATHER_TOOL]), + USER_MSG, + ] + ) + + def test_duplicate_against_request_tools_rejected(self) -> None: + with pytest.raises(ValueError, match="Duplicate tool name"): + self.check( + [dynamic_system_msg([WEATHER_TOOL]), USER_MSG], + tools=[WEATHER_TOOL], + tool_choice="auto", + ) + + +class TestKimiExtensionMapping: + def apply(self, model_type: str = "kimi_k3", **kwargs) -> ChatCompletionRequest: + req = make_request(**kwargs) + _apply_kimi_chat_extensions(req, model_type) + return req + + def kwargs_of(self, req: ChatCompletionRequest) -> dict: + return req.chat_template_kwargs or {} + + def test_non_kimi_untouched(self) -> None: + req = self.apply(model_type="llama", thinking={"type": "disabled"}, top_p=1.0) + assert req.chat_template_kwargs is None + assert req.top_p == 1.0 + + def test_thinking_effort_wins_over_reasoning_effort(self) -> None: + req = self.apply( + thinking={"type": "enabled", "keep": "all", "effort": "low"}, reasoning_effort="max" + ) + assert self.kwargs_of(req)["thinking"] is True + assert self.kwargs_of(req)["thinking_effort"] == "low" + + @pytest.mark.parametrize("effort", ["low", "high", "max"]) + def test_reasoning_effort_applies_when_thinking_effort_absent(self, effort: str) -> None: + req = self.apply(thinking={"type": "enabled", "keep": "all"}, reasoning_effort=effort) + assert self.kwargs_of(req)["thinking_effort"] == effort + + def test_reasoning_effort_none_disables_thinking_when_alone(self) -> None: + req = self.apply(reasoning_effort="none") + assert self.kwargs_of(req)["thinking"] is False + assert "thinking_effort" not in self.kwargs_of(req) + + def test_reasoning_effort_none_does_not_override_explicit_thinking(self) -> None: + req = self.apply(thinking={"type": "enabled", "keep": "all"}, reasoning_effort="none") + assert self.kwargs_of(req)["thinking"] is True + assert "thinking_effort" not in self.kwargs_of(req) + + def test_no_effort_derived_for_explicitly_disabled_thinking(self) -> None: + req = self.apply(thinking={"type": "disabled"}, reasoning_effort="high") + assert self.kwargs_of(req)["thinking"] is False + assert "thinking_effort" not in self.kwargs_of(req) + + def test_medium_effort_has_no_k3_equivalent(self) -> None: + req = self.apply(reasoning_effort="medium") + assert "thinking_effort" not in self.kwargs_of(req) + assert "thinking" not in self.kwargs_of(req) + + def test_default_reasoning_effort_not_mapped(self) -> None: + req = self.apply() + assert req.chat_template_kwargs is None + + def test_client_chat_template_kwargs_win(self) -> None: + req = self.apply(chat_template_kwargs={"thinking": True}, reasoning_effort="none") + assert self.kwargs_of(req)["thinking"] is True + + def test_stream_options_defaulted_for_streaming(self) -> None: + req = self.apply(stream=True) + assert isinstance(req.stream_options, StreamOptions) + assert req.stream_options.include_usage is True + + def test_stream_options_untouched_for_non_streaming(self) -> None: + assert self.apply().stream_options is None + + def test_tool_choice_required_derived_with_tools(self) -> None: + req = self.apply(tools=[WEATHER_TOOL], tool_choice="required") + assert self.kwargs_of(req)["tool_choice"] == "required" + + def test_tool_choice_none_derived_with_tools(self) -> None: + req = self.apply(tools=[WEATHER_TOOL], tool_choice="none") + assert self.kwargs_of(req)["tool_choice"] == "none" + + def test_tool_choice_required_derived_with_dynamic_only_tools(self) -> None: + req = self.apply( + messages=[dynamic_system_msg([WEATHER_TOOL]), USER_MSG], tool_choice="required" + ) + assert self.kwargs_of(req)["tool_choice"] == "required" + + def test_tool_choice_auto_not_derived(self) -> None: + req = self.apply(tools=[WEATHER_TOOL], tool_choice="auto") + assert "tool_choice" not in self.kwargs_of(req) + + def test_response_format_json_object_derived(self) -> None: + req = self.apply(response_format={"type": "json_object"}) + assert self.kwargs_of(req)["response_format"] == "json_object" + + def test_response_format_json_schema_derived(self) -> None: + schema = {"type": "object", "properties": {"city": {"type": "string"}}} + req = self.apply( + response_format={ + "type": "json_schema", + "json_schema": { + "name": "weather", + "schema": schema, + "strict": True, + }, + } + ) + assert self.kwargs_of(req)["response_format"] == "json_schema" + assert self.kwargs_of(req)["response_schema"] == schema + + @pytest.mark.parametrize( + "wrapper, msg", + [ + ({"schema": {"type": "object"}}, "non-empty"), + ({"name": "", "schema": {"type": "object"}}, "non-empty"), + ({"name": "weather"}, "`schema` object"), + ( + {"name": "weather", "schema": {"type": "object"}, "strict": "yes"}, + "must be a boolean", + ), + ], + ) + def test_response_format_json_schema_wrapper_validation(self, wrapper: dict, msg: str) -> None: + with pytest.raises(ValueError, match=msg): + self.apply(response_format={"type": "json_schema", "json_schema": wrapper}) + + +class TestKimiParamPolicy: + @pytest.fixture(autouse=True) + def _enable_policy(self, monkeypatch: pytest.MonkeyPatch) -> None: + # The policy is opt-in (off by default); these tests exercise its + # enabled semantics as a KVV certification run would. + monkeypatch.setenv("TRTLLM_KIMI_PARAM_POLICY", "1") + + def enforce(self, **kwargs) -> ChatCompletionRequest: + req = make_request(**kwargs) + _enforce_kimi_param_policy(req) + return req + + def test_policy_disabled_by_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TRTLLM_KIMI_PARAM_POLICY", raising=False) + req = self.enforce(top_p=0.8, temperature=2.0, n=2) + # Existing deployments keep today's behavior: no rejection, no + # coercion, unless a certification run opts in with =1. + assert req.top_p == 0.8 + assert req.temperature == 2.0 + assert req.n == 2 + + def test_top_p_none_coerced(self) -> None: + assert self.enforce().top_p == 0.95 + + def test_top_p_one_coerced(self) -> None: + assert self.enforce(top_p=1.0).top_p == 0.95 + + def test_top_p_pinned_value_accepted(self) -> None: + assert self.enforce(top_p=0.95).top_p == 0.95 + + def test_top_p_other_rejected(self) -> None: + with pytest.raises(ValueError, match="top_p is fixed"): + self.enforce(top_p=0.8) + + @pytest.mark.parametrize("temperature", [0.0, 0.6, 1.0]) + def test_temperature_in_range_accepted(self, temperature: float) -> None: + self.enforce(temperature=temperature) + + @pytest.mark.parametrize("temperature", [-0.1, 1.1, 2.0]) + def test_temperature_out_of_range_rejected(self, temperature: float) -> None: + with pytest.raises(ValueError, match="temperature"): + self.enforce(temperature=temperature) + + def test_penalties_rejected(self) -> None: + with pytest.raises(ValueError, match="presence_penalty"): + self.enforce(presence_penalty=0.5) + with pytest.raises(ValueError, match="frequency_penalty"): + self.enforce(frequency_penalty=0.5) + + def test_n_rejected(self) -> None: + with pytest.raises(ValueError, match="n is fixed"): + self.enforce(n=2) + + def test_policy_env_off_switch(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("TRTLLM_KIMI_PARAM_POLICY", "0") + req = self.enforce(top_p=0.8, temperature=2.0, n=2) + # Fully unconstrained: no rejection and no coercion. + assert req.top_p == 0.8 + assert req.temperature == 2.0 + assert req.n == 2 + + +class TestKimiResponseFormatGuidedDecoding: + def test_thinking_mode_builds_triggered_tags_on_response_channel(self) -> None: + from tensorrt_llm.serve.openai_protocol import ResponseFormat + + params = _response_format_to_guided_decoding_params( + ResponseFormat(type="json_object"), + reasoning_parser="kimi_k3", + chat_template_kwargs={"thinking": True}, + ) + stag = json.loads(params.structural_tag) + fmt = stag["format"] + assert fmt["type"] == "triggered_tags" + assert fmt["triggers"] == ["<|open|>response<|sep|>"] + assert fmt["tags"][0]["begin"] == "<|open|>response<|sep|>" + assert fmt["tags"][0]["end"] == "<|close|>response<|sep|>" + assert fmt["stop_after_first"] is True + + def test_non_thinking_mode_returns_raw_grammar(self) -> None: + from tensorrt_llm.serve.openai_protocol import ResponseFormat + + params = _response_format_to_guided_decoding_params( + ResponseFormat(type="json_object"), + reasoning_parser="kimi_k3", + chat_template_kwargs={"thinking": False}, + ) + assert params.structural_tag is None + assert params.json_object is True + + +class TestKimiK3StrictGrammar: + def build(self, monkeypatch: pytest.MonkeyPatch, tools: list, gate: str = "1") -> dict | None: + monkeypatch.setenv("TRTLLM_KIMI_K3_STRICT_TOOL_GRAMMAR", gate) + parser = KimiK3ToolParser() + return parser.build_strict_structural_tag_format( + [ChatCompletionToolsParam.model_validate(t) for t in tools] + ) + + def test_env_gate_off_returns_none(self, monkeypatch: pytest.MonkeyPatch) -> None: + strict = json.loads(json.dumps(WEATHER_TOOL)) + strict["function"]["strict"] = True + assert self.build(monkeypatch, [strict], gate="0") is None + + def test_empty_tools_returns_none(self, monkeypatch: pytest.MonkeyPatch) -> None: + assert self.build(monkeypatch, []) is None + + def test_format_shape(self, monkeypatch: pytest.MonkeyPatch) -> None: + strict = json.loads(json.dumps(WEATHER_TOOL)) + strict["function"]["strict"] = True + loose = { + "type": "function", + "function": {"name": "search", "parameters": {"type": "object"}}, + } + fmt = self.build(monkeypatch, [strict, loose]) + assert fmt["type"] == "triggered_tags" + assert fmt["triggers"] == ["<|open|>tools<|sep|>"] + # The deadlock traps: these must stay False or the grammar forbids + # the think/response text before the section and the message close + # after it. + assert fmt["at_least_one"] is False + assert fmt["stop_after_first"] is False + section = fmt["tags"][0] + assert section["begin"] == "<|open|>tools<|sep|>" + assert section["end"] == "<|close|>tools<|sep|>" + calls = section["content"] + assert calls["type"] == "tags_with_separator" + assert calls["separator"] == "" + assert calls["at_least_one"] is True + strict_tag, loose_tag = calls["tags"] + assert strict_tag["begin"] == '<|open|>call tool="get_weather"' + elements = strict_tag["content"]["elements"] + assert elements[0]["type"] == "regex" + assert elements[1] == { + "type": "const_string", + "value": '<|sep|><|open|>json type="object"<|sep|>', + } + assert elements[2]["type"] == "json_schema" + assert elements[2]["json_schema"] == WEATHER_TOOL["function"]["parameters"] + assert strict_tag["end"] == "<|close|>json<|sep|><|close|>call<|sep|>" + assert loose_tag["content"] == {"type": "any_text"} + assert loose_tag["end"] == "<|close|>call<|sep|>" + + def test_tool_name_attribute_escaping_round_trip(self, monkeypatch: pytest.MonkeyPatch) -> None: + name = 'we"ird&name' + tool = { + "type": "function", + "function": {"name": name, "parameters": {"type": "object"}}, + } + fmt = self.build(monkeypatch, [tool]) + begin = fmt["tags"][0]["content"]["tags"][0]["begin"] + escaped = _escape_attr(name) + # Pin the exact K3 dialect: only '&' and '"' have escaped forms. + assert escaped == "we"ird&name" + assert _escape_attr("ac") == "ac" + assert escaped in begin + assert _unescape_attr(escaped) == name + assert _parse_attrs(f'tool="{escaped}" index="1"') == { + "tool": name, + "index": "1", + } + + def test_angle_bracket_tool_name_skips_grammar(self, monkeypatch: pytest.MonkeyPatch) -> None: + # '<' has no escaped form in the K3 wire format and would break the + # parser's attribute regex; the builder must fall back rather than + # emit a grammar the parser cannot read back. + tool = { + "type": "function", + "function": { + "name": "bad None: + xgrammar = pytest.importorskip("xgrammar") + strict = json.loads(json.dumps(WEATHER_TOOL)) + strict["function"]["strict"] = True + fmt = self.build(monkeypatch, [strict]) + xgrammar.Grammar.from_structural_tag(json.dumps({"type": "structural_tag", "format": fmt}))