[Neuron][Kernel] Support Longer Sequences in NKI-based Flash PagedAttention and Improve Efficiency

[lingfanyu]

Summary

This PR is a follow-up of #11277 . It improves code quality and efficiency, and enables kernel to process larger inputs.

Following things are done in this PR:

• fix 2 tiling issues triggered when seqlen_q (i.e. chunk size in chunked-prefill) is larger than 128 and 512
• get rid of the limit of how many KV cache block each tile can access, currently tile_size / block_size <= 128
• unit tests with larger inputs (e.g. seqlen_q > 128, num_blocks_per_tile > 128)
• skip computation in areas masked out by causal masks
• load kv cache once for query heads sharing the same kv head
• remove unused code and format code in pre-commit

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[liangfu]

LGTM, but feel free to chime in.

[lingfanyu]

Thanks!