composable_kernel

mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-04 21:51:28 +00:00

Files

Yi DING 47a04fda08 [rocm-libraries] ROCm/rocm-libraries#5790 (commit c132b5a)

[CK_TILE] Fix NaN for FMHA BWD When seq_q=0

## Motivation
This PR addresses NaNs in the FMHA backward (dQ/dK/dV) path when the
effective query sequence length for a tile is zero, by ensuring the
per-tile pipelines exit early with zeroed accumulators and by avoiding
an early kernel return that prevented writing out cleared gradients.

## Technical Details
- Add unconditional early-exit in the dK/dV pipelines when
`num_total_loop <= 0` (no work), returning zeroed accumulators.
- Adjust group-mode kernel early-return logic to only return when
**both** `seqlen_q` and `seqlen_k` are zero, allowing blocks to run and
store cleared dK/dV when `seqlen_q == 0`.

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

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

2026-03-27 07:54:53 +00:00

fmha_batch_prefill_kernel.hpp

[rocm-libraries] ROCm/rocm-libraries#4999 (commit 45f6624)

2026-03-05 01:09:12 +00:00

fmha_bwd_kernel.hpp

[rocm-libraries] ROCm/rocm-libraries#5790 (commit c132b5a)

2026-03-27 07:54:53 +00:00

fmha_fwd_appendkv_kernel.hpp

chore(copyright): update copyright header for include directory (#3293 )

2025-11-26 11:00:05 -07:00

fmha_fwd_appendkv_tile_partitioner.hpp

chore(copyright): update copyright header for include directory (#3293 )