Skip to content

[CK_TILE] Update CK and add RDNA build support#178

Open
hyoon1 wants to merge 7 commits intock_improve_mainfrom
hyoon1/fa-ck-rdna
Open

[CK_TILE] Update CK and add RDNA build support#178
hyoon1 wants to merge 7 commits intock_improve_mainfrom
hyoon1/fa-ck-rdna

Conversation

@hyoon1
Copy link

@hyoon1 hyoon1 commented Mar 20, 2026

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

@hyoon1 hyoon1 requested a review from rocking5566 March 20, 2026 19:52
@hyoon1 hyoon1 force-pushed the hyoon1/fa-ck-rdna branch from 78af732 to 4a318f7 Compare March 20, 2026 19:58
@hyoon1 hyoon1 changed the title [CK_TILE] Update CK and add gfx1x build support [CK_TILE] Update CK and add RDNA build support Mar 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant