[CK_TILE] Update CK and add RDNA build support

[hyoon1]

Motivation

Update CK to a gfx1x FMHA-capable version and align FlashAttention CK argument wiring accordingly.

Technical Details

• Update composable_kernel to e5683e2
• Align FlashAttention CK integration with the updated CK interface:
  • adjust CK FMHA call/trait argument wiring to match the updated CK side
  • keep CK-side argument mapping consistent after the CK update
• Add ROCm gfx11/gfx12 build support
• Set CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=0 for gfx11 targets
• Add backward guards for unstable gfx1x paths:
  • gfx11: disable CK backward
  • gfx12: disable deterministic CK backward

Test Plan

pytest tests/test_flash_attn_ck.py

Test Result

No failures on both gfx11 and gfx12

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Updates the ROCm CK backend to a newer composable_kernel revision and extends build/runtime support for RDNA (gfx11/gfx12), including updated FMHA argument wiring and targeted backward guards.

Changes:

• Bump composable_kernel submodule and adjust CK FMHA args to match the updated interface.
• Add gfx11/gfx12 build targeting and default CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT behavior for gfx11 targets.
• Gate/skip CK backward in tests and runtime for unsupported/unstable gfx1x backward paths; add optional “LLC Head Grouping” forward dispatch.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/test_flash_attn_ck.py	Adds gfx11/gfx12 detection helpers and skips/guards for unsupported CK backward on gfx1x.
setup.py	Extends supported GPU_ARCHS, adds target detection / forwarding to CK kernel generator, and adjusts CK_TILE defaults for gfx11.
csrc/flash_attn_ck/mha_varlen_fwd.cpp	Updates CK args and adds optional head-grouped forward dispatch with logging.
csrc/flash_attn_ck/mha_varlen_bwd.cpp	Adds gfx1x backward support checks before allocating/intermediate work.
csrc/flash_attn_ck/mha_fwd.cpp	Updates CK args and adds optional head-grouped forward dispatch with logging.
csrc/flash_attn_ck/mha_bwd.cpp	Adds gfx1x backward support checks before allocating/intermediate work.
csrc/flash_attn_ck/flash_common.hpp	Adds ROCm arch detection helpers and a gfx1x deterministic/backward guard.
csrc/composable_kernel	Updates submodule pointer to the newer CK commit.
README.md	Documents RDNA 3/4 support and current backward limitations.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[rocking5566]

Overall looks good — the RDNA support and backward guards are well-structured. A few items worth addressing beyond what Copilot already flagged (I've replied on those threads separately).

[rocking5566]

LGTM