composable_kernel

mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-17 11:30:02 +00:00

Files

ArthurLiu 2a479f7411 [CK][CK_TILE] Fix FMHA codegen group mode dispatch (#6764 )

## Motivation

FMHA codegen had incorrect dispatch behavior in group mode. Two root
causes:

1. Wrong field names in dispatch conditions — Used batch-mode fields
(seqlen_q, seqlen_k) instead of group-mode fields (max_seqlen_q,
max_seqlen_k), causing wrong kernel selection at runtime on gfx950.
2. Missing kernel variants — Group mode was overly filtered out from
smaller-tile specializations (bwd) and lacked spatial-padding pipeline
variants on gfx950 (fwd).

gfx942 don't support trload pipeline. 

## Technical Details

 fmha_bwd.py:
- max_seq_q_cond and extra_cond now emit t.max_seqlen_q / t.max_seqlen_k
for group mode.
- Relaxed kernel filtering: group mode no longer skips tiles with
max_seq_q != 0.

  fmha_fwd.py:
  - get_bm0_cond emits a.max_seqlen_q for group mode tile-size dispatch.
- Added two qr_async_trload pipeline variants with spatial padding for
gfx950 group mode.
  
## Test Plan
Triggering AITER CI job:

## Submission Checklist

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

2026-04-28 02:14:42 +08:00

__init__.py

chore(copyright): update copyright header for example directory (#3273 )

2025-11-24 18:02:41 -08:00

fmha_batch_prefill.py

[CK_TILE] fix(fmha): support >2GB KV cache in batch prefill via template dispatch (#6653 )

2026-04-24 07:08:41 +08:00

fmha_bwd.py

[CK][CK_TILE] Fix FMHA codegen group mode dispatch (#6764 )

2026-04-28 02:14:42 +08:00

fmha_fwd_appendkv.py

[CK] Use as_posix() instead of str() for paths in fmha_fwd_appendkv.py (#4812 )