[Bug]: V2 scheduler deadlocks when a recompute-paused request's replay exceeds max_num_tokens (chunked prefill off) #21826

System Info

• GPU: NVIDIA GeForce RTX 5090 (SM120), driver 595.99.02, CUDA 13.2; Ubuntu 24.04, Python 3.12, torch 2.14.0+cu130
• TensorRT-LLM 1.4.0rc3 (pip) and 1.4.0rc1, PyTorch backend, KV cache manager V2 (the default), KVCacheV2Scheduler
• Also reproduced with main's scheduler_v2.py (6f6fb3458a) loaded on the 1.4.0rc3 runtime. Since rc3, only Helix, DFlash and disaggregation-diagnostics commits have touched it.
• Model: Qwen/Qwen2.5-0.5B-Instruct. Nothing model-specific is involved; the pool is shrunk with max_tokens.

Who can help?

Owners of the V2 scheduler (tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py).

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

Chunked prefill off, GPU-only KV cache, and a load that needs preemption:

import os, random
os.environ.setdefault("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
from tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.llmapi import KvCacheConfig

llm = LLM(model="Qwen/Qwen2.5-0.5B-Instruct", max_batch_size=64, max_num_tokens=2048,
          max_seq_len=4096, enable_chunked_prefill=False,
          kv_cache_config=KvCacheConfig(max_tokens=8192, host_cache_size=0,
                                        enable_block_reuse=False))
rng = random.Random(0)
prompts = [[rng.randrange(100, 30000) for _ in range(1900)] for _ in range(12)]
sp = SamplingParams(max_tokens=512, temperature=0, ignore_eos=True, detokenize=False)
futs = [llm.generate_async(p, sampling_params=sp) for p in prompts]
done = 0
for f in futs:
    try:
        done += len(f.result(timeout=300).outputs[0].token_ids) == 512
    except Exception as e:
        print(type(e).__name__, str(e)[:200])
print(f"completed {done}/12")

The pool holds about 3 requests of 1,900 + 512 tokens; all 12 are submitted at once.

Expected behavior

All 12 requests complete, more slowly, with preemptions. That's what happens with enable_chunked_prefill=True (12/12 in 3.4 s), or when outputs are capped short enough that prompt + output stays at or below max_num_tokens (max_tokens=140: 12/12).

actual behavior

Within about 1 s, 3 requests complete and 9 fail:

[V2Scheduler] Declaring deadlock after 1000 passes: 0 generation and 9 context request(s) blocked; states={'CONTEXT_INIT': 9}; ... cache_tier_below_gpu=False ...
RuntimeError: V2 scheduler deadlock: 0 generation and 9 context request(s) active but none could be scheduled, suspended or preempted in 1000 consecutive attempts. The KV cache pool is likely exhausted. Configure kv_cache_config.host_cache_size, increase kv_cache_config.max_tokens, or lower max_batch_size.

At the raise the pool isn't exhausted. In a larger instance on 1.4.0rc1, 3,410 of 3,411 blocks were free, and none of the blocked requests held a cache.

additional notes

Root cause.

1. Under KV pressure the scheduler preempts a started request by recompute-pause. The executor resets it with reset_for_recompute(_UNBOUNDED_PAUSE_MAX_INPUT_LEN) (py_executor.py#L11380), which deliberately replays the prompt plus all generated tokens as one context.
2. With chunked prefill off, that replay has to run in one iteration. In the repro, once a 1,900-token prompt has generated more than 148 tokens, its replay exceeds max_num_tokens (2048). Tracing the scheduler shows request 67 paused at 2,174 tokens. _try_schedule_context_full then fails budget.can_fit_tokens() and returns STOP on every pass, however empty the pool is (scheduler_v2.py#L1343-L1347).
3. STOP ends the context phase, so no context queued behind it is tried either. That includes request 68, whose own 1,964-token replay would fit. The running requests finish, nothing schedulable is left, and _detect_deadlock raises with "the KV cache pool is likely exhausted" (L2396).

_is_recompute_pause_candidate (L2461) never checks whether a victim's replay can be scheduled again. It takes chunked prefill off, prompt + output able to exceed max_num_tokens, and a load that preempts. A GPU-only KV cache (host_cache_size=0) makes the last much more likely, because pressure can then only be relieved by preemption.

Fix. I'll open a PR with a scheduler-only fix. Everything in it is a no-op with chunked prefill on:

• A. A request whose replay wouldn't fit in max_num_tokens is not a recompute-pause victim.
• B. With no cache tier below GPU, a new context that could outgrow its replay is admitted only if the worst-case pages of every such request, plus the current pages of the others, fit in the pool. Requests A protects can then always grow by preempting the ones that remain preemptible.
• E. If a pending context can never fit max_num_tokens with chunking off, the deadlock error says so and suggests enabling chunked prefill, instead of pointing at an exhausted pool.

With it, every row above completes on 1.4.0rc3, and on main's scheduler (12/12, 24/24). The existing V2 scheduler unit tests still pass.

I searched existing issues and PRs (V2 scheduler deadlock, recompute pause, reset_for_recompute, max_num_tokens replay) and found nothing covering this.