[Gasoonjia]

Summary

The Muse Glimmer sliding-window RingKVCache was hardcoded to 2 * window (4096 for window 2048), and the CUDA export's max prefill chunk is also 2 * window. After a T-token write, its first query still needs the W − 1 keys before it, so W + T − 1 consecutive positions must stay in the ring at once. With ring = chunk = 4096, from the second chunk on, 2047 of 4096 query rows saw only part of their 2048-key window (about 1024 keys missing on average). This affects any chunk=4096 prefill / NLL on the in-graph CUDA path.

The window is fixed by the model, but the chunk is an export choice. So the ring now takes the chunk as a parameter instead of deriving it from the window:

• RingKVCache(..., max_write_len=None) sizes to window + max_write_len - 1, the same rule as the MLX RingBufferKVCache. max_write_len defaults to window, as on MLX.
• MuseGlimmerModel.set_max_prefill_chunk(n) rebuilds the sliding-window rings and keeps _build_masks consistent with them.
• The CUDA solo export (_export_cuda) and the DFlash CUDA export call it with their existing max prefill (min(max_seq_len - 1, 2 * window), i.e. 4096). That gives a ring of 6143. The chunk size is unchanged.
• inference.py chunks eager prefill by model.max_prefill_chunk, so eager chunks always match the ring.

Unchanged paths:

• The DFlash draft ring uses the default max_write_len=window, which covers its writes (min(window, max_target_prefill)).
• The MLX path and the off-graph CUDA KV path already size rings from max_write.

Test plan

• test_prefill_chunks_see_full_window: for chunk sizes 3, window, and 2 * window, every query attends to exactly min(pos + 1, window) keys. It fails with the old 2 * window ring.
• test_ring_sized_for_max_write_len: checks the buf_size formula.
• test_chunked_prefill_matches_sequential (CUDA) now runs max-size chunks that wrap the ring.
• pytest examples/models/muse-glimmer/tests/{test_pipeline,test_export_dflash,test_checkpoint_context,test_mlx_offgraph}.py: 61 passed, 9 skipped (CPU). CUDA tests not run locally (no GPU on this host).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24979

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 44 Pending

As of commit ec3b0f7 with merge base 1ab77be ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

• ❌ The email address for the commit (ec3b0f7) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.