[None][feat] Add PyTorch reset_prefix_cache API

[milesial]

Description

Following vLLM reset_prefix_cache and SGLang flush_cache, add a python API + HTTP endpoint to reset the local KV cache state.
This is useful during benchmarking to reset the state between runs in a concurrency sweep for example.

Test Coverage

Added unit tests to tests/unittest/llmapi/test_llm.py

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

Summary by CodeRabbit

• New Features

  • Added reset_prefix_cache() method to the LLM API (beta status) for PyTorch backend
  • Added POST /reset_prefix_cache endpoint to OpenAI-compatible server

• Tests

  • Added comprehensive test coverage for prefix cache reset functionality across distributed and single-instance executor configurations

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds a reset_prefix_cache() method to invalidate PyTorch-engine KV prefix-cache reuse state. The implementation spans the executor base, PyTorch LLM API (with beta status), and OpenAI server HTTP endpoint, including full test and API stability coverage.

Changes

Prefix Cache Reset API

Layer / File(s)	Summary
Base executor reset method tensorrt_llm/executor/base_worker.py	BaseWorker.reset_prefix_cache() validates engine support and delegates to self.engine.reset_prefix_cache().
PyTorch LLM API with executor dispatch tensorrt_llm/llmapi/llm.py	_TorchLLM.reset_prefix_cache() (beta) validates encode_only and executor availability, routes through collective_rpc() when available, or calls executor method directly; raises NotImplementedError if unsupported.
OpenAI server HTTP endpoint and handler tensorrt_llm/serve/openai_server.py	POST /reset_prefix_cache route and handler validate generator support, invoke reset, and return HTTP 200 or 501 error response.
Tests and API stability documentation tests/unittest/api_stability/references/llm.yaml, tests/unittest/llmapi/test_llm.py	API stability reference documents reset_prefix_cache as beta. Tests verify LLM dispatch to executor (direct, collective RPC, unsupported), OpenAI endpoint success and error cases.

Sequence Diagram

sequenceDiagram
  participant OpenAIClient
  participant OpenAIServer
  participant TorchLLM
  participant Executor
  participant Engine
  OpenAIClient->>OpenAIServer: POST /reset_prefix_cache
  OpenAIServer->>TorchLLM: reset_prefix_cache()
  alt Collective RPC supported
    TorchLLM->>Executor: collective_rpc("reset_prefix_cache")
  else Direct reset
    TorchLLM->>Executor: reset_prefix_cache()
    Executor->>Engine: reset_prefix_cache()
    Engine-->>Executor: cache invalidated
  end
  Executor-->>TorchLLM: success
  TorchLLM-->>OpenAIServer: return
  OpenAIServer-->>OpenAIClient: HTTP 200

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Suggested labels

api-compatible

Suggested reviewers

• longlee0622

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding a PyTorch reset_prefix_cache API. It is concise, specific, and directly related to the changeset across all modified files.
Description check	✅ Passed	The description explains the motivation (following vLLM/SGLang patterns for cache reset during benchmarking), identifies test coverage, and confirms the PR checklist review. The structure aligns with the template requirements.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/llmapi/test_llm.py (1)

2744-2793: ⚡ Quick win

Add missing negative-path coverage for the new reset API contracts.

Coverage is good for dispatch basics, but it’s still missing two high-value guard/error cases in tests/unittest/llmapi/test_llm.py:

1. _TorchLLM.reset_prefix_cache() when llm._encode_only = True (should reject).
2. OpenAIServer.reset_prefix_cache() when server.generator.reset_prefix_cache() exists but raises NotImplementedError (should map to 501, if that is the handler contract).

Suggested test additions

+def test_llm_reset_prefix_cache_rejects_encode_only() -> None:
+    llm = object.__new__(LLM_torch)
+    llm._encode_only = True
+    llm._executor = _FakeResetExecutor()
+
+    with pytest.raises(NotImplementedError):
+        llm.reset_prefix_cache()
+
+
+class _FakeNotImplementedResetGenerator:
+    def reset_prefix_cache(self):
+        raise NotImplementedError("not supported")
+
+
+def test_openai_reset_prefix_cache_endpoint_maps_not_implemented() -> None:
+    server = object.__new__(OpenAIServer)
+    server.generator = _FakeNotImplementedResetGenerator()
+
+    response = asyncio.run(server.reset_prefix_cache())
+    assert response.status_code == 501

As per coding guidelines, tests/** reviews should explicitly assess whether coverage is sufficient and call out concrete follow-up test files when it is not.

🤖 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_llm.py` around lines 2744 - 2793, Add two
negative-path tests: (1) for LLM_torch.reset_prefix_cache create an instance via
object.__new__(LLM_torch), set llm._encode_only = True and llm._executor to any
executor (e.g., _FakeResetExecutor), call llm.reset_prefix_cache() and assert it
raises the expected rejection (e.g., NotImplementedError or the same error type
used for encode-only rejection); (2) for OpenAIServer.reset_prefix_cache set
server = object.__new__(OpenAIServer) and server.generator to a stub whose
reset_prefix_cache method raises NotImplementedError, call
asyncio.run(server.reset_prefix_cache()) and assert the response.status_code ==
501 to verify the NotImplementedError is mapped to 501. Ensure tests reference
LLM_torch.reset_prefix_cache and OpenAIServer.reset_prefix_cache so they cover
the described guard paths.

🤖 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/serve/openai_server.py`:
- Around line 997-1001: The current exception handler around
reset_prefix_cache() only catches NotImplementedError and ValueError so
RuntimeError from _TorchLLM.reset_prefix_cache() bubbles up as a 500; update the
try/except to also catch RuntimeError and map it to the same handler (return
self._create_not_supported_error(str(e))) so encode-only / no-executor cases are
returned as not-supported errors; locate the call to reset_prefix_cache() in
openai_server.py and extend the except clause to include RuntimeError alongside
NotImplementedError and ValueError.

---

Nitpick comments:
In `@tests/unittest/llmapi/test_llm.py`:
- Around line 2744-2793: Add two negative-path tests: (1) for
LLM_torch.reset_prefix_cache create an instance via object.__new__(LLM_torch),
set llm._encode_only = True and llm._executor to any executor (e.g.,
_FakeResetExecutor), call llm.reset_prefix_cache() and assert it raises the
expected rejection (e.g., NotImplementedError or the same error type used for
encode-only rejection); (2) for OpenAIServer.reset_prefix_cache set server =
object.__new__(OpenAIServer) and server.generator to a stub whose
reset_prefix_cache method raises NotImplementedError, call
asyncio.run(server.reset_prefix_cache()) and assert the response.status_code ==
501 to verify the NotImplementedError is mapped to 501. Ensure tests reference
LLM_torch.reset_prefix_cache and OpenAIServer.reset_prefix_cache so they cover
the described guard paths.

🪄 Autofix (Beta)

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: c27f4417-0459-4abc-9663-4477f3c2151d

📥 Commits

Reviewing files that changed from the base of the PR and between 33b0a32 and 4f30f45.

📒 Files selected for processing (5)

• tensorrt_llm/executor/base_worker.py
• tensorrt_llm/llmapi/llm.py
• tensorrt_llm/serve/openai_server.py
• tests/unittest/api_stability/references/llm.yaml
• tests/unittest/llmapi/test_llm.py

[achartier]

Devin raised a TOCTOU race, could you check it? https://app.devin.ai/review/NVIDIA/TensorRT-LLM/pull/14970

[milesial]

Devin raised a TOCTOU race, could you check it? https://app.devin.ai/review/NVIDIA/TensorRT-LLM/pull/14970

@achartier thanks, implemented Devin's feedback

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #53606 [ run ] triggered by Bot. Commit: f271822 Link to invocation

[tensorrt-cicd]

PR_Github #53606 [ run ] completed with state SUCCESS. Commit: f271822
/LLM/main/L0_MergeRequest_PR pipeline #42752 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[milesial]

/bot run

[tensorrt-cicd]

PR_Github #53642 [ run ] triggered by Bot. Commit: ab45757 Link to invocation

[tensorrt-cicd]

PR_Github #53642 [ run ] completed with state FAILURE. Commit: ab45757
/LLM/main/L0_MergeRequest_PR pipeline #42787 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[milesial]

/bot run

[tensorrt-cicd]

PR_Github #53648 [ run ] triggered by Bot. Commit: ab45757 Link to invocation

[tensorrt-cicd]

PR_Github #53648 [ run ] completed with state SUCCESS. Commit: ab45757
/LLM/main/L0_MergeRequest_PR pipeline #42792 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[milesial]

/bot run

[tensorrt-cicd]

PR_Github #53678 [ run ] triggered by Bot. Commit: dd8936e Link to invocation

[tensorrt-cicd]

PR_Github #53678 [ run ] completed with state SUCCESS. Commit: dd8936e
/LLM/main/L0_MergeRequest_PR pipeline #42818 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[milesial]

/bot skip --comment "Known flaky tests"

[tensorrt-cicd]

PR_Github #53711 [ skip ] triggered by Bot. Commit: dd8936e Link to invocation

[tensorrt-cicd]

PR_Github #53711 [ skip ] completed with state SUCCESS. Commit: dd8936e
Skipping testing for commit dd8936e

Link to invocation

[2ez4bz]

[xxi-nv]

Hi @milesial — heads up: this PR introduced a regression on main that breaks the H100_PCIe-PyTorch-Ray-1 stage, so I opened a revert: #15306.

Root cause: this PR adds reset_prefix_cache() to BaseWorker. Since RayGPUWorker(RpcWorkerMixin, BaseWorker) inherits it, the method now exists on the worker class. The Ray worker-extension injection check (RayGPUWorker._inject_worker_extension, ray_gpu_worker.py:196-203) hard-fails on any extension whose attribute name already exists on the worker class, and the RL/verl WorkerExtension also defines reset_prefix_cache:

ValueError: Worker class RayGPUWorker already defines 'reset_prefix_cache',
which conflicts with extension WorkerExtension.
  -> RuntimeError: RayGPUWorker died during initialization

This kills the Ray worker at init, so all test_llm_update_weights*, test_llm_partial_update_weights*, and async_llm tests in that stage fail before reaching their logic (e.g. post-merge L0_Test-x86_64-Single-GPU #3385).

Note the H100_PCIe-PyTorch-Ray-1 stage was already failing in this PR's final CI run (L0_MergeRequest_PR #42818, commit dd8936e) — that failure wasn't flaky, it was this collision — but the PR was merged with /bot skip --comment "Known flaky tests", so it reached main.

Re-landing suggestion: either pick a non-colliding method name on BaseWorker, or make the extension-injection check tolerate an intentional override instead of raising on every name match. Happy to help review the follow-up.