Fix ChunkedKVCache dropping valid tokens during decode (llama4)

[binaydhakal]

Summary

ChunkedKVCache.maybe_trim_front (used only by Llama 4) decides how much of the
KV cache to trim by comparing self.keys.shape[2] — the padded buffer size,
which is over-allocated in blocks of step (256) — against chunk_size. But
valid tokens are front-filled at [0 : offset - start_position], so once the
padded buffer reaches chunk_size the trim slices off the live tokens and keeps
the unfilled padding, and start_position is advanced by the wrong amount.

After the padded buffer crosses chunk_size, every Llama 4 decode step then
attends over corrupted keys, so incremental generation diverges from a single
full forward pass. The generic model test only checks output shapes, so it does
not catch this. The repo's own Llama 4 test config (attention_chunk_size=8,
cache step=256) triggers it on the first decode step; a real Llama 4
(attention_chunk_size=8192) triggers it on long contexts once the padded buffer
reaches 8192.

Reproduction

With random weights, comparing a single forward pass to token-by-token decode
over a 12-token (multi-chunk) sequence, max |Δlogit|:

	before	after
attention_chunk_size=4	3.4	1.3e-06
attention_chunk_size=8	2.8	1.2e-06

Position 0 (prefill) was already correct; every decoded position was wrong. A
dense control model matches to ~1e-6 both before and after, confirming the
invariant and the harness. Verified across batch > 1, longer sequences, and
several chunk sizes (multiple chunk-boundary crossings).

Fix

Trim on the number of valid tokens (offset - start_position) and keep the
most recent chunk_size valid tokens. Those always contain the current chunk
(its start is >= offset - chunk_size), and the existing block-position mask
excludes any older tokens still in the window.

Test

Adds test_llama4_chunked_kv_cache: builds Llama 4 with a small
attention_chunk_size and asserts that incremental decode matches a single
forward pass over a multi-chunk sequence. It fails on main and passes with the
fix. tests/test_models.py and tests/test_prompt_cache.py remain green.