composable_kernel

mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-19 22:39:03 +00:00

Files

Jeff Huang 518d24e662 Add sequence padding and variable length support in fmha (#2932 )

* * [CK_TILE] Add sequence padding and variable length support in fmha (and v3)

 - Group Mode Padding: Introduces the `-s_qpad` argument to support
   physically padded layouts. Kernels now use padded start pointers
   (`seqstart_padded_*_ptr`) for memory addressing.

 - Batch Mode Variable Length: Adds `-q_eff_lens` and `-kv_eff_lens`
   arguments for efficient processing of variable-length sequences by
   passing cumulative effective lengths (`cu_seqlen_*_ptr`) to the kernel.

 - FMHA examples: Support padding and variable length both in
   group and batch mode. Dispatcher is updated as well (dispatch to
   kPadSeqLenK enabled pipeline).

 - New padding test cases: Add padding test cases to `smoke_test_fwd.sh` and
   `test_fmha_fwd.inc`, and add benchmarks to `benchmark_fwd.sh` and
   `benchmark_fwd_v3.sh` as well. These test cases and benchmarks that
   specifically validate/benchmark the new padding and variable-length
   functionalities in both group and batch modes.

* [CK_TILE] Fix build error in fmha unit tests

* [CK_TILE] add mqa, gqa to sequence padding unit tests

* [CI_TILE] Reduce the number of padding seqlen unit tests in FMHA to avoid timeouts in CI

* [CK_TILE] remove unnecessary MageKArgs overload in FmhaFwdV3Kernel and FmhaFwdKernel

2025-09-26 12:36:27 +08:00

fmha_batch_prefill_kernel.hpp

Fix for Add the API to load SGPR (#2913 )

2025-09-25 10:32:42 -07:00

fmha_bwd_kernel.hpp

Fix for Add the API to load SGPR (#2913 )

2025-09-25 10:32:42 -07:00

fmha_fwd_appendkv_kernel.hpp

Fix for Add the API to load SGPR (#2913 )

2025-09-25 10:32:42 -07:00

fmha_fwd_appendkv_tile_partitioner.hpp

[CK_TILE] Add PagedAttention kernels (#1387 )