Fix SinkCache RoPE coordinates for absolute generation

#2
by jholl-hugging - opened

Problem

In Transformers 4.52.1, the standard generation loop advances Qwen2 RoPE
positions absolutely. The upstream SinkCache truncates its KV window but also
rerotates every surviving key into a bounded local slot. Its stored key phase
therefore no longer matches the new query's absolute phase.

Change

  • The custom generate() wrapper selects absolute-position mode. A shift
    retains each surviving key's original RoPE phase while dropping old K/V.
  • Direct callers can select position_mode="local" when both position_ids
    and cache_position intentionally use bounded slots. That path computes
    rerotation in FP32, normalizes the derived twiddle, then writes BF16 keys.
  • Appending exactly to capacity no longer evicts a token early.
  • README documents the position contract and the tested Transformers version.

Verification

  • Qwen2.5-1.5B-Instruct, BF16, eager attention, Transformers 4.52.1, RTX 4080;
    two cache arms interleaved against the same frozen model and WikiText-2
    teacher-forced token stream. Prefill logits matched bitwise.
  • 2,600 tokens, four sinks, 256 recent positions, absolute positions:
    2,339 evictions; upstream 65,492 rerotations and 27.840 post-fill PPL;
    corrected cache zero rerotations and 19.227 PPL. Every reported block
    favored the correction.
  • On the identical model/data/tokens/prefill with bounded local positions,
    correctly rerotating the keys gave 19.268 PPL upstream and 19.236 with
    FP32 normalized rerotation. The corrected absolute and correctly local
    contracts therefore agree closely.
  • At a 1,024-recent-token window and 20,000 tokens with bounded local
    positions, upstream and FP32 normalized rerotation were nearly tied:
    10.122 versus 10.136 post-fill PPL. The numerical change improves an
    isolated retained key's fidelity but is not a demonstrated next-token
    quality gain under that local contract.
  • Eleven regression tests pass under Transformers 4.52.1, including native
    Qwen2 generation, exact capacity fill, and repeated BF16 key shifts.
  • A paired 20,000-token, 1,024-recent-token absolute-position follow-up is
    queued on the same GPU; I will add its result to this PR when it completes.

Scope

The paired model-forward test reproduces the relevant Qwen2 generation
position/cache contract while keeping the input tokens identical between
arms. It is not a throughput benchmark. This historical custom cache is not
compatible with the changed Cache API in Transformers 5.x; the README now
states the tested version.

Validation update (26 September 2026): the paired, paper-sized absolute-position run completed on beast.

  • Qwen2.5-1.5B-Instruct, BF16, native eager attention, Transformers 4.52.1; four sinks, 1,024 recent tokens, 20,000 WikiText-2 raw train tokens. Both arms used the same frozen model and token IDs. The teacher-forced model.forward loop passed monotonically increasing position_ids and cache_position, matching the relevant generation position contract.
  • Before eviction, the two arms' prefill logits had the same SHA-256. Each arm then made 18,971 one-token evictions. RoPE cosine/sine values reached every one of the 531,216 native cache updates. The upstream arm rerotated retained keys 531,188 times across 28 layers; the corrected arm rerotated them zero times, as required for absolute positions.
  • Post-fill perplexity: 14.932650 upstream β†’ 10.017987 corrected. The correction won all 19 reported blocks. Whole-stream perplexity was 15.006433 β†’ 10.276187. The saved result records every block and the model, dataset, token, and source hashes (job-38071dde859b; baseline source SHA-256 ac8d2f4f..., corrected source SHA-256 71ddeb23...).
  • A separate W=1,024 bounded-local-position run on the same train token stream was nearly tied: 10.121977 upstream vs 10.135523 with FP32 normalized rerotation. The substantial absolute-position improvement therefore tracks the coordinate mismatch repaired here, rather than the numerical FP32 change. The exact-capacity append change was not exercised in either full comparison because prefill began at capacity.

The result is a native-model, teacher-forced cache comparison, not a sampled end-to-end model.generate() quality benchmark or a Transformers 5.x compatibility test. An independent WikiText-2 test split run (job-0f39060a3150) is queued; I will add its result here when it finishes.

Independent WikiText-2 test-split check completed (job-0f39060a3150). This used the same Qwen2.5-1.5B-Instruct model, BF16 eager attention, Transformers 4.52.1, four sinks, 1,024 recent tokens, and the exact baseline and corrected source hashes from the 20,000-token train run. The test parquet and token IDs have different hashes from train.

Over 10,000 test tokens, the paired teacher-forced native model.forward loop yielded post-fill perplexity 15.982478 upstream β†’ 10.440735 corrected (whole-stream 14.647943 β†’ 9.997112). The correction won all nine reported blocks. Both arms had identical prefill-logit hashes, 8,971 evictions, valid absolute positions, and RoPE cosine/sine delivery on all 251,216 native cache updates. Upstream rerotated retained keys 251,188 times across 28 layers; the corrected arm made zero rerotation calls.

Baseline generate.py SHA-256: ac8d2f4f7ae25656989b0d7d359790fcb0d04fdb4735265a90f6d4570fc5ed52; corrected SHA-256: 71ddeb239961648ef3ab82a81da3a53a9a0d918941d3cf2f287f69aae4d1f5c5. The test dataset SHA-256 is 5f1bea067869d04849c0f975a2b29c4ff47d867f484f5010ea5e861eab246d91, and the token-ID SHA-256 is 9ef56bb298258264dc0ac20ac0c98bb9a0d60e914c790467259d05714033b272.

This replicates the absolute-position cache repair on an independent text split. It remains a teacher-forced comparison of the generation position/cache contract, not a sampled end-to-end model.generate() benchmark, Transformers 5.x compatibility result, or evidence of a unique integer-phase advantage. The separate bounded-local-position comparison remains near-tied and does not reproduce the paper's much larger manual BF16-arithmetic loss.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment