[FMHA] Batch Prefill Support Improvements: Change KV Cache Layout & Large Page Size Support

[Jeff-Huang]

Proposed changes

This PR introduces significant improvements to the Batch Prefill kernels in ck_tile, enabling codegen support, change the KV layouts and support 1024 page size.

Codegen Support: Enabled code generation for batch_prefill kernels.
* Updated fmha_batch_prefill.py to generate kernels with different page sizes 1024.
* Pipeline & Problem Structure:
* Introduced BlockFmhaBatchPrefillPipelineProblem structure for better management of batch prefill problems.
* Updated block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp to handle async operations and offsets
correctly.

• KV Cache & Layout Enhancements:
  • Added support for k cache layout: [num_blocks, num_kv_heads, head_size/8, block_size, 8].
  • Added support for v cache layout: [num_blocks, num_kv_heads, block_size/8, head_size, 8].
  • Added kv_last_page_lens to kernel arguments to support handling varying sequence lengths in paged attention.
  • Refined kv_offset_array_transform to calculate k/v offset.
  • Added support for large page size (1024).
• Bug Fixes:
  • Fixed incorrect page ID calculation in kv_offset_array_transform specifically for gfx950.

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

• I have added tests relevant to the introduced functionality, and the unit tests are passing locally
• I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
• I have added inline documentation which enables the maintainers with understanding the motivation
• I have removed the stale documentation which is no longer relevant after this pull request
• (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
• I have run clang-format on all changed files
• Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

[afagaj]

@Jeff-Huang This is a significant PR.

Please add one or more entries to CHANGELOG.md as appropriate.

[Jeff-Huang]

@Jeff-Huang This is a significant PR.

Please add one or more entries to CHANGELOG.md as appropriate.

Thanks for the reminder! I've updated the CHANGELOG.md with the relevant entries in the latest commit.