[Bug]: KVCacheV2Scheduler retains active PEFT adapters after KV suspension
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
System Info
- TensorRT-LLM branch:
main - Exact commit:
e53a70c0e3322483e83d5d976b3fe6352e0e022a - Component: PyTorch backend,
KVCacheV2Scheduler, PEFT/LoRA request lifecycle - Reproduction type: source-level scheduler state-machine harness with mocked KV/PEFT managers; no model weights or external data
- Host used for the harness: macOS 26.6.2, arm64, Python 3.12.14
- GPU/CUDA/TensorRT: not used by this controller-level reproduction
Who can help?
Leaving this unassigned for the core maintainers to route.
Information
This was reproduced with a minimal scheduler harness rather than an official example script.
Tasks
No model or dataset is required; the issue is in request/resource lifecycle handling.
Reproduction
With the PyTorch backend using KVCacheV2Scheduler and a PEFT device cache that can hold one adapter:
- Run a generation request using adapter A and queue a context request using adapter B.
- Make KV V2 generation allocation fail under
MAX_UTILIZATION, causing request A to be self-evicted (the same problem also applies to victim eviction). _suspend_request()suspends A in the main and draft KV managers, but does not callmark_request_done(A, pause=True)on the PEFT manager.- Adapter A therefore remains active on device. When the scheduler admits request B, PEFT
ensure_batchcannot evict A's pages from a full device cache.
The missing transition is also documented by the current TODO in scheduler_v2.py.
A minimal mocked state-machine harness asserts the two lifecycle transitions needed for a complete fix:
suspend: peft_cache_manager.mark_request_done(request_a, pause=True)
resume: peft_cache_manager.add_request_peft(request_a, True)
On the exact main commit above, the harness fails both assertions:
AssertionError: suspend: Expected 'mark_request_done' to have been called once. Called 0 times.
AssertionError: resume: Expected 'add_request_peft' to have been called once. Called 0 times.
The resume assertion models the state after A has been correctly marked paused. Generation resumption returns directly from _try_schedule_generation() and does not pass through the context admission path that normally re-adds PEFT ownership.
Expected behavior
KV V2 suspension should release the suspended request's PEFT device ownership while retaining its paused/host-cache state. If that LoRA generation request later obtains KV pages and resumes, its PEFT active-request ownership should be restored before it is scheduled.
This should let a different adapter use a full PEFT device cache without losing the suspended request's ability to resume.
actual behavior
KV V2 suspension only releases KV pages. The suspended LoRA request remains active in PeftCacheManager, so its device adapter pages are not evictable. This can prevent another adapter from being loaded even though the scheduler evicted the original request to relieve resource pressure.
additional notes
I propose a focused two-transition fix:
- Call
mark_request_done(req, pause=True)when KV V2 suspends a request. - When a previously suspended LoRA generation request successfully obtains KV pages, call
add_request_peft(req, True)before returning it as scheduled.
The hot path for requests without PEFT/LoRA remains unchanged. I have regression tests for both the device-page release and generation-resume ownership restoration. Against the exact commit above, the source-level A/B harness fails on the base and passes with the focused patch.
I have not yet run the native TensorRT-LLM test suite or a supported-GPU end-to-end test, so I am not making a performance claim. I searched open issues and PRs for the KV V2 PEFT suspend/resume transitions and found no active duplicate. Per CONTRIBUTING.md, I will submit the focused PR after TensorRT-LLM engineers approve this issue request.
- I searched relevant issues and checked the documentation and examples before submitting.
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 in tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py at the TODO around _suspend_request() and the generation-resume path in _try_schedule_generation(). Run the mocked KV/PEFT state-machine harness described in the issue, then verify that suspension releases PEFT ownership and resumed generation restores it. Done means both lifecycle assertions pass without changing the non-PEFT hot path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai, backend
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100