NVIDIA / NVIDIA/TensorRT-LLM

[Bug]: Scaffolding role-filtered messages_to_retain may produce an unintended KV cache prefix boundary

Open
#19,319 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Scaffolding
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, calls releaseSubtree(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:

  1. Record the target token sequence and the role-filtered retained token sequence.
  2. Compare the computed retention count with the token boundary required by the agreed policy.
  3. Test boundaries around KV block sizes and verify which subtree is detached.
  4. Include single-turn, multi-turn, and tool-use cases, and template-added tokens.
  5. 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.