NVIDIA / NVIDIA/TensorRT-LLM

[Bug]: KVCacheV2Scheduler retains active PEFT adapters after KV suspension

Open
#18,407 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

KV-Cache Management Lora/P-tuning Pytorch
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:

  1. Run a generation request using adapter A and queue a context request using adapter B.
  2. Make KV V2 generation allocation fail under MAX_UTILIZATION, causing request A to be self-evicted (the same problem also applies to victim eviction).
  3. _suspend_request() suspends A in the main and draft KV managers, but does not call mark_request_done(A, pause=True) on the PEFT manager.
  4. Adapter A therefore remains active on device. When the scheduler admits request B, PEFT ensure_batch cannot 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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.