[#17146][fix] Resolve reasoning mode from the rendered prompt - #17305
[#17146][fix] Resolve reasoning mode from the rendered prompt#17305joerowell wants to merge 1 commit into
Conversation
47118cf to
616da4e
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe change adds prompt-based thinking-state resolution, Poolside V1 parser support, Qwen3.5 and MiniMax M2 registrations, Laguna alias mapping, server integration, and expanded tests. ChangesReasoning parser flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ChatTemplate
participant OpenAIServer
participant ReasoningParserFactory
participant PoolsideV1ReasoningParser
ChatTemplate->>OpenAIServer: Render chat prompt
OpenAIServer->>ReasoningParserFactory: Resolve prefilled thinking state
ReasoningParserFactory->>PoolsideV1ReasoningParser: Check prompt markers
PoolsideV1ReasoningParser-->>ReasoningParserFactory: Return thinking state
ReasoningParserFactory-->>OpenAIServer: Return true, false, or None
OpenAIServer->>OpenAIServer: Update thinking and enable_thinking
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
tests/unittest/llmapi/test_reasoning_parser.py (1)
370-374: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicated opt-in test.
test_resolve_prefilled_thinking_opted_inrepeatstest_alias_resolves_identicallywithparser="poolside_v1". The prompt, tails, and expectations are identical, andtest_resolve_prefilled_thinkingalso covers both tails. The negative opt-in case is covered separately bytest_resolve_prefilled_thinking_requires_opt_in.♻️ Proposed removal
-@pytest.mark.parametrize(("tail", "expected"), [(R1_START, True), - (R1_END, False)]) -def test_resolve_prefilled_thinking_opted_in(tail: str, expected): - assert ReasoningParserFactory.resolve_prefilled_thinking( - "poolside_v1", f"<assistant>{tail}") is expected - -🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_reasoning_parser.py` around lines 370 - 374, Remove the redundant test_resolve_prefilled_thinking_opted_in test, retaining test_alias_resolves_identically, test_resolve_prefilled_thinking, and test_resolve_prefilled_thinking_requires_opt_in as the existing coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/llmapi/reasoning_parser.py`:
- Around line 96-97: Add a return annotation to the class method
ReasoningParserFactory.keys, using the appropriate mapping view type for the
object returned by cls._parsers.keys(). Preserve the existing implementation and
behavior.
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 1571-1585: Update the request-processing flow around
ReasoningParserFactory.resolve_prefilled_thinking so the prompt-resolved
thinking and enable_thinking values are applied before
add_thinking_budget_logits_processor runs. Ensure the budget processor for
poolside_v1 and laguna receives the resolved mode rather than stale request
kwargs, while preserving consistent parser selection during postprocessing.
In `@tests/unittest/llmapi/test_reasoning_parser.py`:
- Around line 400-413: Update test_resolve_prefilled_thinking_requires_opt_in to
first assert each parser name is registered in ReasoningParserFactory._parsers,
then retain the existing None assertions for every tail value. This ensures the
test validates opt-in behavior rather than passing for an unregistered name.
- Around line 377-397: Update test_resolved_mode_overrides_stale_thinking_kwarg
to include a scenario where only enable_thinking is set to False while stale
thinking=True remains, and assert that the parser places the visible text in
reasoning_content rather than content. Retain the existing case verifying that
clearing both keys yields normal content, so the test covers the parser’s OR
behavior and proves both keys must be cleared.
---
Nitpick comments:
In `@tests/unittest/llmapi/test_reasoning_parser.py`:
- Around line 370-374: Remove the redundant
test_resolve_prefilled_thinking_opted_in test, retaining
test_alias_resolves_identically, test_resolve_prefilled_thinking, and
test_resolve_prefilled_thinking_requires_opt_in as the existing coverage.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 80888a38-a585-4a62-af3e-642c43dc6998
📒 Files selected for processing (7)
docs/source/developer-guide/telemetry.mdtensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/reasoning_parser.pytensorrt_llm/serve/openai_server.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/api_stability/references/trtllm_serve_cli.yamltests/unittest/llmapi/test_reasoning_parser.py
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
616da4e to
b62ee09
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unittest/llmapi/test_reasoning_parser.py (2)
416-432: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the non-opted-in parser list from the registry.
The hardcoded list does not cover parsers registered later. If a new parser sets
resolves_thinking_from_promptby mistake, this test still passes. Compute the list fromReasoningParserFactory.keys()minus the opted-in names. The registration guard at Line 429 then becomes unnecessary.♻️ Proposed refactor
-@pytest.mark.parametrize("parser", [ - "deepseek-r1", "deepseek_v4", "qwen3", "qwen3_5", "minimax_m2", - "minimax_m3", "nemotron-v3", "nano-v3", "gemma4", "kimi_k2", "kimi_k25" -]) +_OPTED_IN_PARSERS = {"poolside_v1", "laguna"} + + +@pytest.mark.parametrize( + "parser", + sorted(set(ReasoningParserFactory.keys()) - _OPTED_IN_PARSERS)) def test_resolve_prefilled_thinking_requires_opt_in(parser: str): """Parsers that have not opted in must never be resolved from the prompt. `deepseek_v4` shares the base class and `nemotron-v3` / `nano-v3` also read `enable_thinking`, so without the flag they would silently pick up a mode the server inferred. """ - # Otherwise a typo or a dropped registration passes vacuously, since an - # unknown name also resolves to None. - assert parser in ReasoningParserFactory.keys() for tail in (R1_START, R1_END, ""): assert ReasoningParserFactory.resolve_prefilled_thinking( parser, f"<assistant>{tail}") is None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_reasoning_parser.py` around lines 416 - 432, Update test_resolve_prefilled_thinking_requires_opt_in to derive its parameterized parser list from ReasoningParserFactory.keys(), excluding the registered parsers that explicitly opt in via resolves_thinking_from_prompt. Remove the redundant registry-membership assertion, while preserving the existing assertions that each non-opted-in parser resolves to None for all tested prompt tails.
370-374: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the redundant opt-in test.
test_resolve_prefilled_thinking_opted_inrepeats assertions that already exist. Lines 353-354 coverpoolside_v1with both markers, andtest_alias_resolves_identicallycovers the same two tails forpoolside_v1andlaguna. This test adds no new coverage.♻️ Proposed removal
-@pytest.mark.parametrize(("tail", "expected"), [(R1_START, True), - (R1_END, False)]) -def test_resolve_prefilled_thinking_opted_in(tail: str, expected): - assert ReasoningParserFactory.resolve_prefilled_thinking( - "poolside_v1", f"<assistant>{tail}") is expected - -🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/test_reasoning_parser.py` around lines 370 - 374, Remove the redundant test_resolve_prefilled_thinking_opted_in test and its parametrized cases; retain the existing coverage for poolside_v1 markers and alias behavior in the surrounding tests.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/llmapi/test_reasoning_parser.py`:
- Around line 283-457: Add finalization coverage to
test_poolside_v1_reasoning_parser_stream for an unterminated </think> prefix
such as “reason</th”. Feed the fragment through parse_delta, then call finish()
and assert the buffered suffix is emitted as reasoning content rather than
dropped, with no visible content.
---
Nitpick comments:
In `@tests/unittest/llmapi/test_reasoning_parser.py`:
- Around line 416-432: Update test_resolve_prefilled_thinking_requires_opt_in to
derive its parameterized parser list from ReasoningParserFactory.keys(),
excluding the registered parsers that explicitly opt in via
resolves_thinking_from_prompt. Remove the redundant registry-membership
assertion, while preserving the existing assertions that each non-opted-in
parser resolves to None for all tested prompt tails.
- Around line 370-374: Remove the redundant
test_resolve_prefilled_thinking_opted_in test and its parametrized cases; retain
the existing coverage for poolside_v1 markers and alias behavior in the
surrounding tests.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 821cb581-67c1-439f-ade6-410365d6bc21
📒 Files selected for processing (7)
docs/source/developer-guide/telemetry.mdtensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/reasoning_parser.pytensorrt_llm/serve/openai_server.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/api_stability/references/trtllm_serve_cli.yamltests/unittest/llmapi/test_reasoning_parser.py
🚧 Files skipped from review as they are similar to previous changes (6)
- docs/source/developer-guide/telemetry.md
- tests/unittest/api_stability/references/trtllm_serve_cli.yaml
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/serve/openai_server.py
- tensorrt_llm/llmapi/reasoning_parser.py
DomBrown
added
the
api-compatible
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Approach looks right — reading the mode off the rendered prompt is the only thing that works when the template prefills the marker, and gating it on resolves_thinking_from_prompt with a test that every other registered parser stays at None is the right way to keep it contained. The rename matches what tool_parser_factory.py:38 already does, so doing it here is fine.
One ask before merge: the parsers that opt into resolves_thinking_from_prompt are only correct on the code path that can actually see a rendered prompt (openai_chat with add_generation_prompt). On every other path — offline LLM API, and the disagg generation server, which receives prompt_token_ids — laguna/poolside_v1 now silently falls back to IdentityReasoningParser and returns reasoning text as content. That is a behavior regression relative to today, and it fails silently: no error, no log line, just a response with the wrong field populated. Please add a guard so an unsupported path is visible rather than quietly wrong. Concretely:
- A one-time
logger.warningwhen a parser withresolves_thinking_from_prompt = Trueis constructed without a resolvable mode (nothinking/enable_thinkinginchat_template_kwargs), naming the parser and saying the mode could not be resolved so output will not be split. - A startup-time check in the disagg generation server: if
reasoning_parseropts into prompt resolution and the server role is generation, that combination can never work today. I'd lean toward a hard error there rather than a warning, since it's a static config mistake that produces wrong output on every request, but a loud startup warning is acceptable if you'd rather not break existing deployments in a bugfix PR.
I wouldn't hard-fail the offline LLM path — callers that pass chat_template_kwargs themselves are legitimate and shouldn't be blocked — so a warning is the right level there.
Two gaps also worth a line in the PR description: the laguna key is not behavior-preserving, and the fix doesn't reach disaggregated serving. Details inline.
| dtype=np.int32).tolist() | ||
|
|
||
| rendered_prompt = None | ||
| if request.prompt_token_ids is not None: |
There was a problem hiding this comment.
In disaggregated serving the generation server receives prompt_token_ids (relayed b64 from the context server, decoded a few lines above) and skips async_apply_chat_template entirely, so rendered_prompt stays None — but the gen server is the one that runs apply_reasoning_parser and produces the user-visible response. The fix therefore doesn't apply under disagg, and combined with the base-class change it's a step backwards there: laguna used to at least split on an emitted </think>, now it will return everything as content.
Unlike the offline case, this one is a static configuration property, not a per-request one: if the server role is generation and reasoning_parser is a parser with resolves_thinking_from_prompt = True, every request will be parsed in the wrong mode, and nothing in the response indicates it. That's worth a gate at startup rather than a note in the release notes. My preference would be to raise during server init on that combination (fail fast on a config that cannot produce correct output), with a logger.warning as the fallback if you'd rather not turn an existing deployment into a hard failure inside a bugfix PR — but silently degrading is the one option I'd rule out.
The real fix — ctx server relays the resolved mode, or the rendered prompt is propagated — is reasonable to defer to a follow-up; the gate is what keeps the gap from being invisible in the meantime.
There was a problem hiding this comment.
Went with @zhaoyangwang-nvidia's suggestion below instead of gating. DisaggregatedParams is already relayed wholesale from ctx to gen, so it just needed one more field: ctx stamps resolved_thinking, gen uses it when nothing was rendered.
So it's fixed rather than gated. I haven't run it against a live disagg deployment though, only unit tests either side, so that hop could use another pair of eyes.
There was a problem hiding this comment.
Relay approach is better than gating, agreed. I traced the hop instead of deploying it and it holds: _get_ctx_request pins stream: False (openai_disagg_service.py:214), so the ctx worker always lands in chat_response_post_processor — the one handler you stamp — and _get_gen_request copies ctx_response.choices[0].disaggregated_params wholesale (openai_disagg_service.py:231), overwriting only request_type/schedule_style/conversation_id/ctx_usage. So the field survives end to end.
Two things left. The relay isn't gated on the parser opting in — separate inline comment on that. And the correctness of the streaming handler not stamping depends entirely on that forced stream: False, which is unstated. Rather than a comment next to the streaming handler, I'd put your original gate on the consuming side instead: gen worker, opted-in parser, nothing rendered, and no relayed value → once-per-process logger.warning. That one condition catches every route into the silently-wrong mode (a future non-_get_ctx_request orchestration path, a rolling upgrade where the ctx worker predates the field, a hand-crafted generation_only request), not just the streaming one, and it's the same if as the opt-in gate. Written up on the openai_server.py inline comment.
| # Templates that prefill <think>/</think> leave the marker in the | ||
| # prompt, so the request kwargs alone cannot tell the parser which | ||
| # mode was rendered. Take it from the prompt instead. | ||
| if (postproc_args.reasoning_parser and rendered_prompt |
There was a problem hiding this comment.
This wiring has no test — test_poolside_v1_mode_resolved_from_prompt hand-simulates the two calls, so nothing catches it if this block gets dropped, if add_generation_prompt gating changes, or if postproc_args.chat_template_kwargs stops being what apply_reasoning_parser reads (postprocess_handlers.py:165). A small test under tests/unittest/llmapi/apps/ with a stub template that prefills <think> / </think> would pin the end-to-end path cheaply.
Note also that request.add_generation_prompt being false lands in the same silent-fallback bucket as the offline and disagg paths: the condition just doesn't fire and the parser runs in whatever mode the request kwargs implied. Whatever warning you add for the unresolved case should cover this branch too, so all three unsupported paths report the same way.
Also note add_thinking_budget_logits_processor (line 1483) still builds its parser from the raw request.chat_template_kwargs, before rendering. It happens to work here because reasoning_start/reasoning_end are class attributes on DeepSeekV4ReasoningParser, so _get_reasoning_boundary finds them even on the Identity branch — worth a comment so the coupling isn't accidental.
There was a problem hiding this comment.
Both done. New test is tests/unittest/llmapi/apps/test_reasoning_prompt_resolution.py, in l0_cpu.yml, driving the real ChatPostprocArgs and apply_reasoning_parser against a stub prefilling template.
It's CPU-only, so it mirrors the server block rather than importing it. Catches the chat_template_kwargs contract breaking, won't catch the block being deleted. That needs a GPU test with a custom --chat_template, which I'd rather leave for a follow-up.
Comment added on the budget coupling.
There was a problem hiding this comment.
The new file covers the chat_template_kwargs contract, which was the main thing I wanted pinned, and deferring the custom---chat_template GPU test is fine. One follow-up in the same file though: the two disagg tests (test_context_worker_stamps_the_resolved_mode, test_generation_worker_uses_the_relayed_mode) re-implement the relay inline rather than calling it, so they pass by construction. The ctx half is reachable from CPU — call chat_response_post_processor with an output carrying disaggregated_params and assert resolved_thinking on the resulting choice. Details in the inline comment.
b62ee09 to
2a6bce9
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/llmapi/test_reasoning_parser.py`:
- Around line 288-294: Add complete type annotations to all 14 new
reasoning-parser test functions in
tests/unittest/llmapi/test_reasoning_parser.py at lines 288-294, 305-316,
326-333, 336-341, 347-349, 360-362, 365-367, 372-374, 384-399, 407-415, 418-434,
445-454, 461-473, and 481-496: add -> None return annotations, use dict[str,
bool] for kwargs, list[str] for stream fixtures, and bool or bool | None for
expected values as applicable.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 50931dfd-1840-4f75-8af0-fea07f415824
📒 Files selected for processing (7)
docs/source/developer-guide/telemetry.mdtensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/reasoning_parser.pytensorrt_llm/serve/openai_server.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/api_stability/references/trtllm_serve_cli.yamltests/unittest/llmapi/test_reasoning_parser.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tests/unittest/api_stability/references/trtllm_serve_cli.yaml
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/serve/openai_server.py
- docs/source/developer-guide/telemetry.md
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Reviewed the change. The approach (resolving the rendered mode from the prompt, behind an opt-in flag) is sound and the unit tests are thorough. What blocks me is the laguna semantics change: it now resolves to a kwargs-driven parser, while the fix only covers the "OpenAI chat endpoint + server-rendered prompt + add_generation_prompt=True" path. Offline LLM API users and the disagg generation server regress from "splits on an emitted </think>" to "everything is content" — as @brnguyen2 already pointed out; I checked and I agree those are real. Either give those paths an actual fix, or keep the old behavior and warn loudly when the mode cannot be resolved, rather than silently degrading. The rest below are non-blocking.
| base64.b64decode(request.prompt_token_ids_b64), | ||
| dtype=np.int32).tolist() | ||
|
|
||
| rendered_prompt = None |
There was a problem hiding this comment.
Following up on @brnguyen2's disagg point: one option beyond gating at startup is that the context server does render the prompt, so it could resolve the mode there and relay it to the generation server alongside prompt_token_ids_b64, instead of the gen server having to infer it. Either way, could the unresolved cases (thinking is None, add_generation_prompt=False, non-str prompt from a processor) log a warning rather than silently falling back to the request kwargs?
There was a problem hiding this comment.
Done, this was the better idea. DisaggregatedParams already gets copied onto the gen request, so ctx stamps resolved_thinking and gen reads it.
On the warning: nothing degrades any more, those paths now behave as they do on main, so it'd be informational rather than a signal. Say if you still want it.
| parser_cls = entry[0] | ||
| if not parser_cls.resolves_thinking_from_prompt: | ||
| return None | ||
| tail = prompt.rstrip() |
There was a problem hiding this comment.
prompt.rstrip() allocates a copy of the whole prompt on every request, but only the last few characters matter. prompt[-64:].rstrip() gives the same result at constant cost.
There was a problem hiding this comment.
Done, prompt[-64:].rstrip() behind a constant.
| return None | ||
| tail = prompt.rstrip() | ||
| end = getattr(parser_cls, "reasoning_end", None) | ||
| start = getattr(parser_cls, "reasoning_start", None) |
There was a problem hiding this comment.
A parser that opted into resolves_thinking_from_prompt necessarily defines reasoning_start/reasoning_end, so the getattr(..., None) fallback defends an unreachable state — and it turns a typo into a silent None (i.e. a silent fallback to the request kwargs). Suggest declaring both on the base that carries this contract and accessing them directly.
There was a problem hiding this comment.
Good catch. The silent None fell back to the request kwargs, which is the exact failure this PR is meant to remove. Both markers are ClassVar[str] on the base now and accessed directly, so opting in without defining them raises.
| # Opt in on parsers whose template prefills the reasoning marker into the | ||
| # prompt and that select their mode from `enable_thinking`. Only those can | ||
| # have the mode resolved from the rendered prompt. | ||
| resolves_thinking_from_prompt = False |
There was a problem hiding this comment.
Nit: annotate as ClassVar[bool] to make it explicit that this is a class-level switch, not an instance field.
| return self._parser.finish() | ||
|
|
||
|
|
||
| @register_reasoning_parser("poolside_v1") |
There was a problem hiding this comment.
Nit: register_reasoning_parser takes *keys, and both keys share the same defaults here, so @register_reasoning_parser("poolside_v1", "laguna") on one line reads better than two stacked decorators (the stacked form elsewhere exists because each layer passes different kwargs).
Also: if laguna is meant to be a compatibility alias only, consider a deprecation warning when it is used — right now it is indistinguishable from poolside_v1 in the CLI choices, the telemetry allowlist and the manifest, so users have no signal to migrate.
There was a problem hiding this comment.
Collapsed.
On deprecation, I'd rather not. laguna shipped in an RC and I'm keeping it as a permanent alias, not a transitional one, so there's nothing to migrate to. poolside_v1 is what our checkpoints declare and what vLLM and SGLang use, so that's the name I'd point people at, but the old one shouldn't start warning.
2a6bce9 to
7000db3
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
7000db3 to
5fafc13
Compare
e76115b to
a81dd4d
Compare
Laguna templates prefill <think> or </think>, so the model never emits the opening marker and the request kwargs cannot tell the parser which mode was rendered. Read it off the prompt instead, behind an opt-in flag so only poolside_v1 is affected. Adds poolside_v1, keeping laguna as an alias. Signed-off-by: Joe Rowell <joerowell4@gmail.com>
a81dd4d to
1348236
Compare
|
Thanks both for the reviews. @zhaoyangwang-nvidia I think yours was read against an earlier state. Disagg is properly fixed now too, following your suggestion, so the mode is relayed from the context worker rather than inferred by the generation worker. Minor thing: on main Rebased onto current main as of this morning, which also picks up #17157. Could you take another look? |
|
/bot run |
DomBrown
left a comment
DomBrown
left a comment
There was a problem hiding this comment.
Approving from API perspective
|
PR_Github #64606 [ run ] triggered by Bot. Commit: |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Re-reviewed after the three fixes. The unresolved-mode fallback in PoolsideV1ReasoningParser.__init__ restores the pre-PR laguna behavior exactly, so the base-class swap no longer changes any path that isn't newly resolved — that was my main concern and it's settled.
I traced the disagg relay rather than deploying it, and the hop holds up: _get_ctx_request forces stream: False (openai_disagg_service.py:214), so the ctx worker always lands in chat_response_post_processor — the one handler you stamp — and _get_gen_request copies ctx_response.choices[0].disaggregated_params wholesale (openai_disagg_service.py:231) before overwriting only request_type/schedule_style/conversation_id/ctx_usage. So the field survives.
That relay does rest on an unstated dependency, though: the streaming post-processor deliberately does not stamp, which is correct only because of the forced stream: False. Rather than a comment there, I'd put a once-per-process logger.warning on the consuming side — gen worker, nothing rendered, opted-in parser, and no relayed value — since that single condition catches every route into the silently-wrong mode (a future orchestration path that doesn't go through _get_ctx_request, a rolling upgrade where the ctx worker predates the field, a hand-crafted generation_only request), not just the streaming one. Details in the [openai_server.py:1585](https://github.com/NVIDIA/TensorRT-LLM/pull/17305/files#diff-2eb783dbd719ea8b2067105279f5a50460168065912d0efd2a3b2cdc0b3a78e5R1585) comment; it's the same if as the opt-in gate.
Remaining comments are narrow — the relay isn't gated on the parser opting in, two of the new tests re-implement the server logic rather than calling it (so they pass by construction), and resolve_prefilled_thinking's docstring under-specifies both why a suffix test is sound and what its three return values mean.
Description hygiene: the PR body still describes only the prompt-resolution flag and the rename. The new DisaggregatedParams.resolved_thinking field and the ctx→gen relay are a protocol addition that reviewers of a disagg deployment will care about; please add a line. The rename mirrors tool_parser/poolside_v1_parser.py and its "laguna" -> "poolside_v1" alias map, so it's consistent with what's already there — no complaint. No user-facing docs mention laguna, and the telemetry table/manifest/CLI reference are all updated, so nothing else is owed on docs.
|
|
||
| @classmethod | ||
| def keys(cls): | ||
| def resolve_prefilled_thinking(cls, reasoning_parser: str, |
There was a problem hiding this comment.
Two suggestions for this helper's contract, both doc/naming fixes.
Document why a suffix test is correct, not just cheap. The docstring says templates "prefill the marker" but not where. In these templates the marker is the last thing appended after the assistant header — ...<|assistant|><think> when thinking is on, ...<|assistant|></think> when off (the stub in tests/unittest/llmapi/apps/test_reasoning_prompt_resolution.py matches this). The two prefills are mutually exclusive and both land at the very end, so which marker ends the prompt is the mode: a trailing unclosed <think> means the template opened reasoning; a trailing </think> means it already closed it and all model output is content.
The suffix restriction is what makes this correct, not merely fast: prior assistant turns render with their own <think>...</think> pairs, so a contains check would misfire on every multi-turn request. The docstring should state that.
One more line worth adding: _PROMPT_TAIL_CHARS is a copy bound, not a semantic window — whitespace before the marker is unbounded, but more than ~56 chars of trailing whitespace after it would push the marker out of the slice and read as unresolved.
Keep the tri-state return, but fix the name/docstring. Optional[bool] reads smelly, but the three cases are genuinely distinct at the call site (openai_server.py:1574-1585):
True/False— the template resolved the mode, overriding the request.None— this prompt can't answer; the caller falls through to the disagg relay, then to the parser's own unresolved fallback.
Folding None into False would turn "couldn't tell" into "explicitly thinking-off", suppressing the relay and silently picking a mode.
That's why is_using_prefilled_thinking is the wrong name: an is_ prefix promises a bool, and None becomes a trap for anyone writing if not is_using_prefilled_thinking(...). Either keep the current name (it reads as a resolution, which is what it does) or rename to something like resolve_thinking_from_rendered_prompt. Either way, spell out the contract:
Returns:
True - the template prefilled `<think>`: reasoning is open.
False - the template prefilled `</think>`: reasoning is already
closed, so all model output is content.
None - the mode cannot be determined from this prompt (unknown
parser, parser has not opted in, or neither marker is at
the tail). Callers must treat this as "ask elsewhere"
(e.g. the relayed disagg value), not as thinking-off.
| if disaggregated_params is not None and args.chat_template_kwargs: | ||
| # Relay the mode we resolved from the rendered prompt; the | ||
| # generation worker never renders and so cannot resolve it. | ||
| resolved = args.chat_template_kwargs.get("enable_thinking") |
There was a problem hiding this comment.
This stamps for any reasoning parser, not just the ones that opted into prompt resolution — args.chat_template_kwargs["enable_thinking"] here is just whatever the client sent. Combined with the gen-side block, which writes both thinking and enable_thinking from the relayed value, that changes behavior for unrelated parsers: a request with {"thinking": true, "enable_thinking": false} against deepseek_v4 gets True in aggregated serving (the parser ORs the two keys) but False under disagg, because the relay flattens both keys to the enable_thinking value.
Degenerate input, but the fix is one line on each side — only stamp/consume when the configured parser has resolves_thinking_from_prompt, which is exactly the set this feature is meant for.
| if thinking is None and request.disaggregated_params is not None: | ||
| # Generation worker: it never rendered, so use the mode the | ||
| # context worker resolved and relayed. | ||
| thinking = request.disaggregated_params.resolved_thinking |
There was a problem hiding this comment.
Two asks on this branch, both the same if.
Gate it on the parser opting in. As written this consumes resolved_thinking for whatever parser is configured, and then overwrites both keys of the caller's chat_template_kwargs (see the ctx-side comment on [postprocess_handlers.py:422](https://github.com/NVIDIA/TensorRT-LLM/pull/17305/files#diff-f18f66c51e14fc204af740548177c93251297a4d1e0ab8c9f00fbb08a79b8ad8R422) for the concrete misparse). Suggest only entering the disagg branch when ReasoningParserFactory.resolves_thinking_from_prompt(postproc_args.reasoning_parser) — a small classmethod would keep the registry lookup in one place, alongside resolve_prefilled_thinking.
Warn when the mode is unresolvable here. Once that gate exists, this is also the only place in the system with enough information to notice the bad state: opted-in parser, rendered_prompt is None (so nothing to resolve locally), and request.disaggregated_params.resolved_thinking is None. Every request on such a deployment is then parsed in a possibly-wrong mode — in thinking mode the prompt prefilled <think>, so the model's output has no opening tag, the fallback DeepSeekR1Parser(reasoning_at_start=False) waits for a marker that never arrives, and reasoning text lands in content with nothing in the response saying so. Same class of wrong answer this PR set out to fix.
A once-per-process logger.warning on that condition covers every route into it, not just one: a future orchestration path that doesn't go through _get_ctx_request (which is what currently forces stream: False and thereby guarantees the ctx worker hits the stamping handler), a rolling upgrade where the ctx worker is an older build without the field, or a client hand-crafting a generation_only request. I'd keep it a warning rather than a hard error — a mixed-version disagg upgrade would otherwise take out the whole gen tier, and the fallback still returns a usable, if unsplit, response.
I'd prefer this over a sentinel stamped from chat_stream_post_processor: ChatCompletionResponseStreamChoice has no disaggregated_params field at all ([openai_protocol.py:806](https://github.com/NVIDIA/TensorRT-LLM/pull/17305/files#diff-4b6ece0a77c7a186737d731752536269c760e4aa3e9ebc896f8023a37280d10aR806)-812) and the streaming handler never calls to_disaggregated_params, so a sentinel there is dead code whose only live path requires someone to first add disagg-param plumbing to the streaming choice — at which point they're already editing the code that would carry the stamp. It would also mean widening resolved_thinking past Optional[bool] on a wire type with extra="forbid", for a state that can't currently occur.
|
|
||
|
|
||
| @pytest.mark.parametrize("relayed", [True, False]) | ||
| def test_generation_worker_uses_the_relayed_mode(relayed: bool) -> None: |
There was a problem hiding this comment.
This test and test_context_worker_stamps_the_resolved_mode above re-implement the two server halves inline (thinking = params.resolved_thinking, then assert the parser behaves accordingly), so they pass by construction and can't fail if the ctx stamp in [postprocess_handlers.py:419](https://github.com/NVIDIA/TensorRT-LLM/pull/17305/files#diff-f18f66c51e14fc204af740548177c93251297a4d1e0ab8c9f00fbb08a79b8ad8R419) or the gen block in [openai_server.py:1583](https://github.com/NVIDIA/TensorRT-LLM/pull/17305/files#diff-2eb783dbd719ea8b2067105279f5a50460168065912d0efd2a3b2cdc0b3a78e5R1583) changes or is dropped. That's the one hop you flagged as unverified against a live deployment, so it's the place where a construction-only test is least useful.
The ctx side at least is reachable from CPU: build a fake GenerationResult-ish output with disaggregated_params set and call chat_response_post_processor directly, then assert response.choices[0].disaggregated_params.resolved_thinking is True. That pins the stamp itself rather than a copy of it.
| return None | ||
| # Only the tail matters, so avoid copying the whole prompt. | ||
| tail = prompt[-_PROMPT_TAIL_CHARS:].rstrip() | ||
| if tail.endswith(parser_cls.reasoning_end): |
There was a problem hiding this comment.
parser_cls.reasoning_start / .reasoning_end are declared on BaseReasoningParser as bare ClassVar[str] annotations with no values, so they only exist on classes that assign them. DeepSeekR1Parser sets them in __init__ as instance attributes — a future parser that subclasses it and flips resolves_thinking_from_prompt = True would hit AttributeError here, in the request path, not at startup.
Cheap insurance: read them with getattr(parser_cls, "reasoning_end", None) and return None (or raise at registration time) when either is missing, so the failure is a startup error rather than a per-request 500.
|
PR_Github #64606 [ run ] completed with state
|
6 participants
Laguna templates prefill
<think>or</think>, so the model never emits the opening marker and the request kwargs cannot tell the parser which mode was rendered. Thus, we should read it from the prompt instead.To make this change minimal, I've kept it behind an opt-in flag so only poolside_v1 is affected.
As an aside, I've renamed the laguna parser to poolside_v1, to keep it consistent with the rest of the world. Happy to defer this to a follow-up PR, but figured I'd do it in one go. We keep the alias, so it shouldn't break anything.
Dev Engineer Review
<think>or</think>markers.poolside_v1and retainedlagunaas a compatibility alias.qwen3_5,minimax_m2, andminimax_m2_append_thinkregistrations.thinkingandenable_thinking.QA Engineer Review
tests/unittest/llmapi/test_reasoning_parser.py.tests/integration/test_lists/entries were modified.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.