kv-cells : scan the flat pos array instead of walking the used std::set

[abdel-darwish-27]

Overview

The kv cache token lookup was walking a std::set every token which showed up as ~17% of cpu at long context. Changed it to scan the flat pos array instead, same cells same order. Decode at 130k went 17 to 19.7 tok/s, outputs identical.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - used AI to help write the code, profiled and tested it myself on a dual RTX A6000 box

[ggml-gh-bot]

Hi @abdel-darwish-27, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.