NVIDIA / NVIDIA/TensorRT-LLM

KV cache connector hangs with Eagle speculative decoding when rewind crosses block boundary

Open
#16,448 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

KV-Cache Management Speculative Decoding
Dominant language
Python
Stars
14.7k
Forks
2.8k
Avg merge
2d 23h
Merged PRs (30d)
489

Description

Description

When using a KV cache connector with Eagle speculative decoding, the executor can hang during get_finished. The hang occurs when rejected draft tokens cause rewindKVCache to free a block that crosses a block boundary.

Root cause

KvCacheConnectorSchedulerOutputRequest.block_ids only grows monotonically — it is never trimmed after rewind. When rewindKVCache calls releaseLastBlock (which happens when the rewind crosses a block boundary), the connector retains stale references to the now-freed block ids.

This causes get_finished's cross-rank mpi_allgather + set.intersection to never complete, because the freed block ids are still present in the connector's per-request bookkeeping.

Repro

  • Backend: PyTorch
  • Spec decoding: Eagle (non-linear tree)
  • KV cache connector: enabled (e.g. KVBM)
  • Observe: hang after the first iteration where Eagle rejects draft tokens and the rewind crosses a block boundary

Suggested fix

Add an on_rewind hook to the connector that trims block_ids (and related per-request state) to match the post-rewind live cache indices after each rewindKVCache call.

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 by tracing rewindKVCache, releaseLastBlock, and get_finished, then inspect KvCacheConnectorSchedulerOutputRequest.block_ids and its per-request bookkeeping. Reproduce the Eagle speculative-decoding case with a KV cache connector enabled and observe the cross-rank hang. Done means rewinding across a block boundary leaves only live block references and get_finished completes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.