[Bug]: Scaffolding role-filtered messages_to_retain may produce an unintended KV cache prefix boundary
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Summary
Scaffolding builds DropKVCacheTask.messages_to_retain by filtering all system/user messages from a conversation, including user messages that occur after assistant/tool messages. The engine, however, uses only the token count of this filtered conversation as a prefix-retention threshold on the target conversation's cache path.
For interleaved multi-turn conversations, these are different semantics: the length of a re-rendered, role-filtered conversation does not necessarily identify the intended prefix boundary in the original conversation.
This issue requests clarification of the intended retention policy and a corresponding fix/test. It does not claim demonstrated inference-output corruption or measured performance loss.
Scope / environment
Source inspected at upstream main commit 7b1bedbbbd949b57725c896c37c97240cf406ca2. Findings below are based on static source inspection; no GPU/model-backed reproduction has been run.
Separate from the HTTP endpoint mismatch tracked by #19281 / #19282. That routing issue must be resolved or bypassed to exercise the message-based truncate path end to end.
Evidence
1. Retained messages are selected by role, not by a contiguous prefix
tensorrt_llm/scaffolding/task.py:438:
self.messages_to_retain = [
message for message in chat_task.messages
if message.role in ("system", "user")
]
For example:
Conversation: system → user1 → assistant1 → tool → assistant2 → user2
Selected: system → user1 → user2
The selected messages are not a contiguous prefix of the conversation. This matters because a ChatTask can accumulate multiple turns before the scope ends.
2. Common construction point and all decorator call sites
The single direct production construction site is in drop_kv_cache_scope(), after normal completion of the decorated controller:
tensorrt_llm/scaffolding/task_collection.py:1139.
Static decorator uses found in the inspected revision:
| Controller | Status | Source |
|---|---|---|
| Coder | Active decorator | tensorrt_llm/scaffolding/contrib/Coder/coder.py:113 |
| MultiroundChatController | Active decorator | examples/scaffolding/benchmarks/multiround_chat_benchmark.py:689 |
| DummyMultiConersation | Test-only, dummy worker | tests/unittest/scaffolding/test_task_collection.py:173 |
| SWEBenchCoder | Commented out | tensorrt_llm/scaffolding/contrib/Coder/coder.py:134 |
| IterResearchController | Commented out | tensorrt_llm/scaffolding/contrib/iter_research/agent.py:81 |
| Researcher | Commented out | tensorrt_llm/scaffolding/contrib/open_deep_research/researcher.py:250 |
| Supervisor | Commented out | tensorrt_llm/scaffolding/contrib/open_deep_research/supervisor.py:61 |
An active decorator does not imply an HTTP request is always sent: drop_kv_cache_handler returns early unless kv_cache_hint_enabled is enabled:
tensorrt_llm/scaffolding/worker.py:450.
The token-based Trace Replay path is separate; it does not use this role-filtering constructor.
3. The worker sends two message lists; the server encodes both
The worker constructs messages from the ChatTask history and messages_to_retain from the filtered list:
tensorrt_llm/scaffolding/worker.py:127.
A qualification: params["messages"] subsequently overrides the initial serialization with chat_task.messages_to_dict_content(), which omits messages whose content is None. Thus this issue does not assume the target serialization always exactly reproduces the original inference prompt:
tensorrt_llm/scaffolding/worker.py:454,
tensorrt_llm/scaffolding/task.py:371.
The server separately applies the chat template and tokenization to both lists:
tensorrt_llm/serve/resource_governor.py:112,
tensorrt_llm/serve/resource_governor.py:150.
The Harmony handler likewise separately converts both lists to token IDs:
tensorrt_llm/serve/resource_governor.py:203.
4. The executor discards retained-token contents and uses only their length
tensorrt_llm/_torch/pyexecutor/py_executor.py:4615:
self.kv_cache_manager.truncate_blocks(
request.messages, len(request.messages_to_retain))
At this point both fields contain token-ID lists, not message objects or strings. The length is therefore a token count, not a character count.
The Python KV manager forwards these arguments unchanged:
tensorrt_llm/_torch/pyexecutor/resource_manager.py:2646.
5. The C++ implementation walks target tokens and applies the numeric threshold
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2221:
- Builds block keys from
targetTokens. - Walks matching blocks using
findMatchingBlock(...). - When
numMatchedTokens > numTokensToKeep, callsreleaseSubtree(matchingBlock).
The filtered retained-token sequence is never used for matching or checked against the target prefix. Actual truncation is block-granular and uses the current strict > condition, not an exact token-level cut.
Expected behavior / policy to clarify
The retention boundary should correspond to an explicitly defined prefix of the actual cached token sequence.
Which behavior is intended at scope completion?
- Retain only the shared system/prompt prefix?
- Retain the initial system + user prefix?
- Retain through a particular later turn, including all preceding assistant/tool tokens?
- Or is role-filtered token count intentionally just a heuristic retention budget?
If the intention is to preserve selected system/user messages while removing intervening assistant/tool messages, prefix truncation cannot implement that operation: the later KV states depend on the intervening context.
Suggested validation
Use a real supported tokenizer/template with at least two user turns separated by assistant/tool messages:
- Record the target token sequence and the role-filtered retained token sequence.
- Compare the computed retention count with the token boundary required by the agreed policy.
- Test boundaries around KV block sizes and verify which subtree is detached.
- Include single-turn, multi-turn, and tool-use cases, and template-added tokens.
- Verify behavior end to end with cache reuse/block evidence after fixing the route mismatch.
No specific replacement policy is proposed yet; clarifying the intended semantics should precede changing the implementation.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with drop_kv_cache_scope in tensorrt_llm/scaffolding/task_collection.py and the messages_to_retain construction in task.py, then trace worker.py, resource_governor.py, and py_executor.py. Review the existing task-collection tests and add multi-turn, tool-use, tokenizer-template, and KV-block boundary coverage after the retention policy is decided. Done means the boundary matches the agreed policy and end-to-end cache reuse is verified.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100