Skip to content

Fix singleton chunks in ChunkPress 🤖🤖🤖 - #325

Open
bingdongni wants to merge 1 commit into
NVIDIA:mainfrom
bingdongni:fix/chunkpress-single-token-tail
Open

bingdongni wants to merge 1 commit into
NVIDIA:mainfrom
bingdongni:fix/chunkpress-single-token-tail

Conversation

@bingdongni

Copy link
Copy Markdown

What changed

Keep a singleton chunk without scoring it. ChunkPress always retains at least one token per chunk, but a final one-token chunk currently reaches scorers such as TOVAPress, which fail when computing attention scores for an empty history. For example, a 9-token prefill with a chunk length of 4 fails on its last chunk.

The new tests cover a single-token input, one-token tails and a chunk length of 1 with TOVA, SnapKV and Knorm, checking both cache length and the retained key/value. They use a small randomly initialized Llama model.

Validation

  • GitHub Actions run / redacted job log: 24 passed (12 CPU, 12 CUDA) on a self-hosted RTX 4060 Laptop GPU, PyTorch 2.6.0+cu124 and transformers 4.57.6.
  • black --check, isort --check-only and flake8 on the two changed files.

Signed-off-by: cuteNBD <206630735+bingdongni@users.noreply.github.com>
@copy-pr-bot

copy-pr-bot Bot commented Nov 7, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant