[None][fix] Kimi-K3: honour the thinking and effort settings of Anthropic and chat requests

tensorrt_llm/serve/extensions/kimi_k3.py

@@ -73,6 +73,27 @@ def enforce_kimi_param_policy(request: ChatCompletionRequest) -> None:
         raise ValueError(f"n is fixed at 1 for this model; got {request.n}.")
 
 
+def _derive_from_reasoning_effort(kwargs: dict[str, Any], effort: object) -> None:
+    """Derive the K3 ``thinking`` / ``thinking_effort`` kwargs from a reasoning effort.
+
+    A no-op when ``kwargs`` already sets an effort or turns thinking off.
+    Otherwise ``"none"`` turns thinking off unless ``kwargs`` sets it,
+    ``"low"``, ``"high"`` and ``"max"`` set the effort, and other values (e.g.
+    harmony's ``"medium"``) have no K3 equivalent and leave the template default.
+    """
+    if (
+        not isinstance(effort, str)
+        or "thinking_effort" in kwargs
+        or kwargs.get("thinking") is False
+    ):
+        return
+    effort = effort.lower()
+    if effort == "none":
+        kwargs.setdefault("thinking", False)
+    elif effort in ("low", "high", "max"):
+        kwargs["thinking_effort"] = effort
+
+
 def dynamic_tool_dicts(messages: Optional[List[ChatCompletionMessageParam]]) -> list[dict]:
     """Collect message-level (dynamic) tool declarations from system messages."""
     tools: list[dict] = []
@@ -124,6 +145,8 @@ def validate_kimi_dynamic_tools(request: ChatCompletionRequest) -> None:
 class KimiK3ServingExtension(ServingExtension):
     """Kimi/Moonshot API semantics for kimi_k3 chat requests."""
 
+    render_version = 2
+
     # Trailing tokens of the native K3 generation prompt (the channel opener
     # ``<|open|>think|response<|sep|>``) that are not reported as prompt usage.
     _GENERATION_CHANNEL_OPENER_TOKENS = 3
@@ -170,7 +193,10 @@ def apply_chat_extensions(self, request: ChatCompletionRequest) -> None:
         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
+        client-supplied ``chat_template_kwargs`` win over derived values.
+        The template reads ``thinking`` / ``thinking_effort``, not the
+        ``enable_thinking`` / ``reasoning_effort`` kwargs most other chat
+        templates read, so those generic kwargs are translated as well. 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.
@@ -183,32 +209,27 @@ def apply_chat_extensions(self, request: ChatCompletionRequest) -> None:
         if request.stream and request.stream_options is None:
             # StreamOptions defaults: include_usage=True, continuous off.
             request.stream_options = StreamOptions()
+        # Sources of the ``thinking`` / ``thinking_effort`` kwargs, highest
+        # precedence first; each one only sets what no source above it set:
+        #   1. those two keys in the client's ``chat_template_kwargs``;
+        #   2. the request-level ``thinking`` object;
+        #   3. the request-level ``reasoning_effort``;
+        #   4. the generic ``enable_thinking`` / ``reasoning_effort`` keys in
+        #      ``chat_template_kwargs`` (the Anthropic adapter maps ``thinking``
+        #      and ``output_config.effort`` onto them);
+        #   5. the template default: thinking on, effort ``max``.
+        # An effort of ``"none"`` turns thinking off, and no effort is derived
+        # for a request whose thinking is off.
         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")
-        ):
-            # 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.
-            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 "reasoning_effort" in request.model_fields_set and request.reasoning_effort is not None:
+            _derive_from_reasoning_effort(
+                derived, getattr(request.reasoning_effort, "value", request.reasoning_effort)
+            )
         if (
             "tool_choice" in request.model_fields_set
             and (request.tools or dynamic_tool_dicts(request.messages))
@@ -234,11 +255,33 @@ def apply_chat_extensions(self, request: ChatCompletionRequest) -> None:
                 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 {}),
-            }
+        chat_template_kwargs = {**derived, **(request.chat_template_kwargs or {})}
+        unset = {"thinking", "thinking_effort"} - chat_template_kwargs.keys()
+        enable_thinking = chat_template_kwargs.get("enable_thinking")
+        if isinstance(enable_thinking, bool):
+            chat_template_kwargs.setdefault("thinking", enable_thinking)
+        _derive_from_reasoning_effort(
+            chat_template_kwargs, chat_template_kwargs.get("reasoning_effort")
+        )
+        translated = unset & chat_template_kwargs.keys()
+        exempted = set(translated)
+        generic_effort = chat_template_kwargs.get("reasoning_effort")
+        if isinstance(generic_effort, str) and generic_effort.lower() in (
+            "none",
+            "low",
+            "high",
+            "max",
+        ):
+            # The extension consumes this control even if a higher-precedence
+            # K3 setting wins. Keep it available to custom templates as well.
+            exempted.add("reasoning_effort")
+        if exempted:
+            # Derived and consumed controls need not be read again by a Jinja template.
+            request.injected_chat_template_kwargs = sorted(
+                {*(request.injected_chat_template_kwargs or ()), *exempted}
+            )
+        if derived or translated:
+            request.chat_template_kwargs = chat_template_kwargs
 
     def structured_output_format(
         self, content: Dict[str, Any], chat_template_kwargs: Optional[Dict[str, Any]]

tests/unittest/llmapi/apps/test_kimi_serve_extensions.py

@@ -388,6 +388,144 @@ def test_response_format_json_schema_wrapper_validation(self, wrapper: dict, msg
             self.apply(response_format={"type": "json_schema", "json_schema": wrapper})
 
 
+def k3_thinking_kwargs(req: ChatCompletionRequest) -> dict:
+    """The K3 template's own thinking controls in the request's merged kwargs."""
+    kwargs = req.chat_template_kwargs or {}
+    return {key: kwargs[key] for key in ("thinking", "thinking_effort") if key in kwargs}
+
+
+class TestKimiGenericThinkingKwargs:
+    """``enable_thinking`` / ``reasoning_effort``: the kwargs other chat templates read."""
+
+    def apply(self, **kwargs) -> ChatCompletionRequest:
+        req = make_request(**kwargs)
+        apply_model_chat_extensions(req, "kimi_k3")
+        return req
+
+    @pytest.mark.parametrize(
+        "generic, expected",
+        [
+            ({"enable_thinking": False}, {"thinking": False}),
+            ({"enable_thinking": True}, {"thinking": True}),
+            ({"reasoning_effort": "low"}, {"thinking_effort": "low"}),
+            ({"reasoning_effort": "high"}, {"thinking_effort": "high"}),
+            ({"reasoning_effort": "max"}, {"thinking_effort": "max"}),
+            ({"reasoning_effort": "High"}, {"thinking_effort": "high"}),
+            ({"reasoning_effort": "none"}, {"thinking": False}),
+            # No K3 equivalent: the template default applies.
+            ({"reasoning_effort": "medium"}, {}),
+            # No effort for a request whose thinking is off ...
+            ({"enable_thinking": False, "reasoning_effort": "high"}, {"thinking": False}),
+            # ... and "none" does not override an explicit thinking setting.
+            ({"enable_thinking": True, "reasoning_effort": "none"}, {"thinking": True}),
+            # Only a bool is a thinking setting.
+            ({"enable_thinking": "false"}, {}),
+        ],
+        ids=[
+            "thinking_off",
+            "thinking_on",
+            "effort_low",
+            "effort_high",
+            "effort_max",
+            "effort_any_case",
+            "effort_none",
+            "effort_medium",
+            "no_effort_when_off",
+            "none_keeps_thinking_on",
+            "non_bool_thinking",
+        ],
+    )
+    def test_generic_kwargs_map_to_the_k3_kwargs(self, generic: dict, expected: dict) -> None:
+        req = self.apply(chat_template_kwargs=generic)
+        assert k3_thinking_kwargs(req) == expected
+
+    def test_k3_kwargs_win_over_the_generic_ones(self) -> None:
+        req = self.apply(
+            chat_template_kwargs={
+                "thinking": True,
+                "thinking_effort": "high",
+                "enable_thinking": False,
+                "reasoning_effort": "low",
+            }
+        )
+        assert k3_thinking_kwargs(req) == {"thinking": True, "thinking_effort": "high"}
+        assert req.injected_chat_template_kwargs == ["reasoning_effort"]
+
+    @pytest.mark.parametrize(
+        "fields, expected",
+        [
+            (
+                {"thinking": {"type": "enabled", "effort": "high"}},
+                {"thinking": True, "thinking_effort": "high"},
+            ),
+            ({"thinking": {"type": "disabled"}}, {"thinking": False}),
+            ({"reasoning_effort": "none"}, {"thinking": False}),
+            ({"reasoning_effort": "high"}, {"thinking": True, "thinking_effort": "high"}),
+        ],
+        ids=["thinking_effort", "thinking_disabled", "reasoning_effort_none", "reasoning_effort"],
+    )
+    def test_request_fields_win_over_the_generic_kwargs(self, fields: dict, expected: dict) -> None:
+        req = self.apply(
+            chat_template_kwargs={"enable_thinking": True, "reasoning_effort": "low"}, **fields
+        )
+        assert k3_thinking_kwargs(req) == expected
+
+    def test_translated_kwargs_are_marked_as_server_injected(self) -> None:
+        # The unused-kwargs guard exempts injected keys when a Jinja template renders.
+        req = self.apply(
+            chat_template_kwargs={"enable_thinking": True, "reasoning_effort": "low"},
+            injected_chat_template_kwargs=["enable_thinking"],
+        )
+        assert req.injected_chat_template_kwargs == [
+            "enable_thinking",
+            "reasoning_effort",
+            "thinking",
+            "thinking_effort",
+        ]
+
+    def test_thinking_off_is_parsed_as_an_answer_from_the_start(self) -> None:
+        from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory
+
+        req = self.apply(chat_template_kwargs={"enable_thinking": False})
+        parser = ReasoningParserFactory.create_reasoning_parser("kimi_k3", req.chat_template_kwargs)
+        result = parser.parse("plain answer<|close|>response<|sep|><|end_of_msg|>")
+        assert (result.content, result.reasoning_content) == ("plain answer", "")
+
+
+class TestKimiAnthropicThinkingControls:
+    """Anthropic ``thinking`` / ``output_config.effort`` reach the K3 template kwargs."""
+
+    def convert(self, **fields) -> ChatCompletionRequest:
+        from tensorrt_llm.serve.anthropic_adapter import convert_anthropic_request
+        from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest
+
+        req = convert_anthropic_request(
+            AnthropicMessagesRequest(model="m", max_tokens=4096, messages=[USER_MSG], **fields)
+        )
+        apply_model_chat_extensions(req, "kimi_k3")
+        return req
+
+    def test_thinking_disabled_turns_k3_thinking_off(self) -> None:
+        req = self.convert(thinking={"type": "disabled"}, output_config={"effort": "high"})
+        assert k3_thinking_kwargs(req) == {"thinking": False}
+        assert "thinking" in req.injected_chat_template_kwargs
+
+    @pytest.mark.parametrize("effort", ["low", "high", "max"])
+    def test_effort_sets_the_k3_effort(self, effort: str) -> None:
+        req = self.convert(
+            thinking={"type": "enabled", "budget_tokens": 1024}, output_config={"effort": effort}
+        )
+        assert k3_thinking_kwargs(req) == {"thinking": True, "thinking_effort": effort}
+        assert "thinking_effort" in req.injected_chat_template_kwargs
+
+    def test_applying_the_extension_again_changes_nothing(self) -> None:
+        # A router renders a request and forwards it; the worker applies the extension again.
+        req = self.convert(output_config={"effort": "low"})
+        kwargs, injected = dict(req.chat_template_kwargs), list(req.injected_chat_template_kwargs)
+        apply_model_chat_extensions(req, "kimi_k3")
+        assert (req.chat_template_kwargs, req.injected_chat_template_kwargs) == (kwargs, injected)
+
+
 class TestKimiParamPolicy:
     @pytest.fixture(autouse=True)
     def _enable_policy(self, monkeypatch: pytest.MonkeyPatch) -> None:

tests/unittest/llmapi/apps/test_render_core.py

@@ -561,6 +561,73 @@ def test_a_gpt_oss_model_without_the_harmony_switch_uses_the_template(
             != render_chat(chat_request(), harmony_resources()).token_ids
         )
 
+    @pytest.mark.parametrize(
+        ("kwargs", "template", "expected"),
+        [
+            ({"reasoning_effort": "LOW"}, "{{ thinking_effort }}", "low"),
+            ({"reasoning_effort": "none"}, "{{ thinking }}", "False"),
+            (
+                {"reasoning_effort": "low", "thinking_effort": "high"},
+                "{{ thinking_effort }}",
+                "high",
+            ),
+            (
+                {"reasoning_effort": "low"},
+                "{{ reasoning_effort }}:{{ thinking_effort }}",
+                "low:low",
+            ),
+        ],
+    )
+    def test_kimi_generic_effort_is_consumed_before_jinja_validation(
+        self, tokenizer, monkeypatch, kwargs, template, expected
+    ) -> None:
+        monkeypatch.delenv("TRTLLM_ALLOW_UNUSED_CHAT_TEMPLATE_KWARGS", raising=False)
+        request = chat_request(chat_template=template, chat_template_kwargs=kwargs)
+        res = RenderResources.from_tokenizer(tokenizer, model_type="kimi_k3")
+
+        assert render_chat(request, res, tokenize=False).text == expected
+
+    def test_kimi_unrecognized_generic_effort_still_requires_a_template_consumer(
+        self, tokenizer, monkeypatch
+    ) -> None:
+        from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError
+
+        monkeypatch.delenv("TRTLLM_ALLOW_UNUSED_CHAT_TEMPLATE_KWARGS", raising=False)
+        request = chat_request(
+            chat_template="{{ thinking_effort | default('max') }}",
+            chat_template_kwargs={"reasoning_effort": "medium"},
+        )
+        res = RenderResources.from_tokenizer(tokenizer, model_type="kimi_k3")
+
+        with pytest.raises(UnusedChatTemplateKwargsError, match="reasoning_effort"):
+            render_chat(request, res, tokenize=False)
+
+    def test_the_generic_thinking_kwargs_reach_the_kimi_k3_prompt(self) -> None:
+        # K3's template reads `thinking` / `thinking_effort`; the extension derives them from the
+        # `enable_thinking` / `reasoning_effort` that other templates read and the Anthropic
+        # adapter sends.
+        res = RenderResources.from_tokenizer(make_tokenizer(name=KIMI_K3), model_type="kimi_k3")
+
+        def render(**chat_template_kwargs) -> str:
+            request = chat_request(
+                messages=[{"role": "user", "content": "hi"}],
+                chat_template_kwargs=chat_template_kwargs or None,
+            )
+            return render_chat(request, res, tokenize=False).text
+
+        default = render()
+        assert default.startswith('<|open|>message role="system" type="thinking-effort"<|sep|>')
+        assert "`thinking_effort=max`" in default
+        assert default.endswith("<|open|>think<|sep|>")
+
+        # Thinking off: the prompt opens with the user turn and ends in the answer channel.
+        off = render(enable_thinking=False)
+        assert off.startswith('<|open|>message role="user"<|sep|>hi')
+        assert off.endswith("<|open|>response<|sep|>")
+
+        low = render(reasoning_effort="low")
+        assert low == default.replace("`thinking_effort=max`", "`thinking_effort=low`")
+
 
 class TestRenderConversation:
     def test_matches_render_chat_for_a_parsed_conversation(self, tokenizer) -> None: