[sebxt]

PR description

ChunkPress sliced hidden_states, keys and values per chunk but forwarded the full kwargs["position_embeddings"] to the inner press. Presses that recompute window queries (SnapKVPress, PyramidKVPress, TOVAPress, CompactorPress, NonCausalAttnPress, DropKVPress) take the last positions of those embeddings, so each chunk's queries were rotated as if they were at the end of the full context. The relative distance to the chunk keys was wrong for every chunk but the last, and the tokens kept in a chunk depended on how much text came after it.

This passes a copy of kwargs with position_embeddings sliced to the chunk. Presses that do not use position_embeddings are unaffected.

The new test compresses the same random hidden states and keys once as a single 32 token chunk and once as the first chunk of a 128 token input, and checks that the first chunk keeps the same keys. It fails before the fix and passes after. tests/presses, tests/test_press_call.py and tests/test_per_layer_compression_press.py pass (732 passed, 4 skipped on CPU).

Checklist

Before submitting a PR, please make sure:

• Tests are working (make test)
• Code is formatted correctly (make style, on errors try fix with make format)
• Copyright header is included
• All commits are signed-off using git commit -s

Fixes #308

AI tools were used to find this bug and to write the fix, the test and this description. The tests were run locally before and after the change.

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

[maxjeblick]

/ok to test 4c373e3

[maxjeblick]

Thanks a lot for the PR, code LGTM