Remap every position in get_freqs, not just the chunk start

[pssrawat]

Summary:
RopeWithAttentionSink.get_freqs has two branches. The static-shape
branch remaps every position and indexes the table with the result. The
dynamic-shape branch remapped only the first position of the chunk and then
took a contiguous slice from there:

remapped_item = sink_size + (input_pos_item - sink_size) % ring_size
freqs_cos = self.freqs_cos.narrow(0, remapped_item, seq_len)

narrow hands row remapped_item + j to the j-th token, so the chunk is
assigned remap(start) + j where it should be remap(start + j).

At sink_size=4, window_size=8 the ring top is 20. A chunk of 5 starting at
18 covers true positions 18..22 and needs rows 18, 19, 4, 5, 6. The slice
asks for 18, 19, 20, 21, 22. Rows 20..22 are outside the ring the cache
actually uses, so those tokens get frequencies for positions no cache slot
holds.

The two branches therefore disagree with each other. On master, the same
wrapping chunk returns rows 18, 19, 4, 5, 6 under static shapes and raises
under dynamic ones, so the frequencies depend on enable_dynamic_shape.

Differential Revision: D116138610

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 98ea915 with merge base 3dcc374 ():

NEW FAILURES - The following jobs have failed:

• Cadence Build & Test / hifi-build / hifi4 (gh)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.
• Cadence Build & Test / vision-build / vision (gh)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.

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

[meta-codesync]

@pssrawat has exported this pull request. If you are a Meta employee, you can view the originating Diff in D116138610.

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