[0z5a]

Summary

ChunkKVPress computes chunk scores for every example but previously applied the first example's selected chunks to the entire batch. Select and gather chunks independently for each example so a later example keeps its own high-scoring tokens.

If selection of a trailing partial chunk gives examples different retained lengths, raise a clear error: the dense cache tensor cannot represent that ragged result. Single-example partial chunks and batches with equal retained lengths continue to work.

Fixes #317.

Validation

• Four CPU tests passed with PyTorch and Transformers 5.2, covering independent selections, partial chunks, ragged-length rejection, and single-example behavior. Three fail on the original implementation.
• make style passed, including flake8, mypy, and SPDX checks; Black and isort checks passed for changed files.
• These are cache-compression tests with deterministic scores; no model accuracy or GPU performance benchmark was run.

[copy-pr-bot]

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.