Files
composable_kernel/include/ck_tile
Anton Gorenko 57d63b3e70 [CK_TILE] Add gtests for FMHA (#2744)
* Improve random number generation

* use different seed for each input (Q, K, V...);
* use deterministic generation of:
  * seqstart_q/k (for group mode);
  * block_table (for paged-kvcahe);
  * cache_batch_idx (for kvcache);

* Extract arg_parser-related code from run functions to use them as tests

* Split examples into main programs and fmha runners, build instances separately

* Add dummy tests that use instances and runners

* Fix a missed corner case of f32->f8 conversion

When value if < min f8 denormal but > min f8 denormal / 2, it must be
rounded to min f8 denormal (i.e. 0b1), not to 0.

* Fix incorrect fp8 scales for P and O in validation code

DataTypeConfig was incorrectly compared with fp8_t.

* Add host generation of dropout random values and use it for validation

Previously host validation (reference_batched_dropout) used random
numbers generated by BlockDropout of the kernel, meaning that incorrect
generation on device (bad distribution, repeated numbers, too many zeros,
etc.) would not trigger any validation errors.

* Implement tests from smoke_test_bwd.sh

* Return result as enum to distinguish failure and missing instance

* Add tests for bwd features: bias, alibi, dropout

* Implement tests from smoke_test_fwd.sh

* Pass seqlen_q/k as vectors to fwd and bwd runners

* Add tests for fwd features: bias, alibi, dropout

* Add tests for pagedkv and splitkv

* Fix conditions when to use splitkv and pagedkv kernels

splitkv was executed only when use_kvcache which == (need_append_kvcache || use_cache_batch_idx || 0 < page_block_size).
In the SplitKV tests: the regular fwd kernel was executed if use_cache_batch_idx was not requested even when num_splitkv > 1.
In the AppendKV tests: the pagedkv kernel was executed but it often failed to find an instance.

* Add tests for appendkv

* Use is_v_rowmajor = true because there are no instances with column layout anymore

* Split public and private compile options for instances

Tests and examples need to know only about CK_TILE_FMHA_FWD_*_API.

* Improve parsing validation in bias and mask

* Pass bias as string for consistency with mask

* Catch parsing and other exceptions

* Add bwd test for deterministic flag

* Initialize fp8 tensors (-init=ufq) similarly to uf

* Fix splitkv/pagedkv invocation: use padded sk when seqlen_k_ptr is not null

seqlen_k cannot be used to determine padding when seqlen_k_ptr is
provided. The actual seqlen_k is taken from seqlen_k_ptr[b].
Even seqlen_k values (% bn0 == 0) use padded seqlen_k while seqlen_k_ptr
may contain arbitrary values.
In the example or tests this produces incorrect results with appendkv
(for example, -d=32 -s=1 -s_k=64 -s_knew=7 -vlayout=c -b=8).

* Fix use_pagedkv value when kvcache = true but page_block_size = 0

In this case block_table_ptr is nullptr which is accessed in the kernel.

* Clean up bwd tests

* Unify fwd tests for f16/bf16 and fp8

* Use better explicit instantiation declaration for fmha_bwd<2>

* Use the same seed for all tests, allow to override it with env variable

* Undo clang-format of one irrelevant file

For some reason my local clang-format-18 and the one in CI work differently.

* Do not build instances and tests on unsupported archs

* Build instance libraries as OBJECT library

* CI: Enable sccache for HIP

There are source files with LANGUAGE HIP, they need
-DCMAKE_HIP_COMPILER_LAUNCHER=sccache

* Add tests to REGRESSION_TESTS

* Fix OOB accesses in deterministic bwd due to incorrectly assumed kN0

The runner assumes kN0 = (hdim_q <= 128) ? 128 : 64 but there are
smaller tiles (for tr_load or fp32). This can create too small dq_acc_buf.

* Pass CK_TILE_FMHA_FWD_*_API as INTERFACE compile options

The instances don't actually depend on them, only examples and tests do.
Passing these definitions as INTERFACE allows to change FMHA_FWD_ENABLE_APIS
without recompiling instances that are already in ccache.

* Fix formatting and names

[ROCm/composable_kernel commit: ec006bb8e0]
2025-09-10 08:06:14 +05:00
..
2025-09-10 08:06:14 +05:00
2024-12-12 11:54:03 +08:00

Back to the main page

Composable Kernel Tile

concept

ck_tile provides a programming model with templated abstractions to enable users to implement performance-critical kernels for machine learning workloads. introduces following basic concepts to help users building your own operator

  • tensor coordinate transformation, this is the core concept of layout/index transform abstraction in both compiler time and run time.
  • tile-based programming model, including tile-level api and the concept of distributed tensor.

ck_tile is independently from the old ck, located under /include/ck_tile. You don't need to include anything from old CK, ck_tile has similiar (indeed almost the same) implementations for users to build operators. We will have a transition period to pull everything from old ck into ck_tile, stay tuned.

component

ck_tile is splitted into several componenets including core, host, ops/gemm, ops/fmha... each component you only need to include a single header (e.g #include "ck_tile/core.hpp", #include "ck_tile/ops/fmha.hpp") then you are able to use the function/structure inside (different from old ck)

[core]
ck_tile/core contains all the basic data structure and function to build the kernel, you can only include this header and build your own operators that utilizing all the basic building blocks introduced in ck.

core/container

  • array, store runtime variables with fixed length (tensor index, register buffer, etc...)
  • tuple, same as std::tuple, hold different type of data, and one of the solution to achieve multiple buffer.
  • sequence, compile time integer sequence used to build various internal structures, or to describe tile size
  • other convenient structure build on top of above 3

core/numeric

  • gpu data type like fp16_t, bf16_t, fp8_t... and the conversion between each other
  • constexpr integer similiar to std::integral_constant to be used as compile time integer.
  • math functions and numeric utilities

core/algorithm

  • coordinate transformation system, used to build tensor transform and compile time indexing. This is the core idea introduced in old ck to describe how a tensor is build by several basic transform primitives like merge/unmerge/embed etc... and how we indexing into a ND tensor that finally mapped to 1D memory offset.

core/tensor

  • tensor descriptor, to describe how a ND tensor
  • distributed tensor, describe the storage of this tensor, and the distribution of how a collection of threads collaborately work for this tensor.
  • tile level API, including load_tile, store_tile, shuffle_tile, slice_tile, etc...

[host]
ck_tile/host contains all the host side utilities to launch a kernel, create the device buffer, and some reference implementations. This can be used to create examples (like that under ck_tile example folder) and simple executable to invoke this kernel, so if you only need ck_tile to build your own device library then it's OK to not include this. Based on this, it is recommended to include the specific header you needed under this folder to avoid including unwanted headers (e.g, only include ck_tile/host/kernel_launch.hpp), unless you are writing a host executable.

[ops/gemm, ops/fmha, ops/reduce...]
our implementation of different device operators.

  • warp, warp tile level operator
  • block, block tile level operator
  • pipeline, pipeline that can achieve a customized tile level mainloop (or epilogue). By switching different pipeline to the kernel template you can have different kind of pipeline optimizations.
  • kernel, template interface for users to instantiate a particular kernel

[ops/epilogue]
epilogue part of our kernel. We may extend this epilogue part to let users to build their own cutomized epilogues.

[ref]
reference implementation of cpu or gpu. This folder is supposed to include a specific header on demand.

examples

currently we put all ck_tile related example under /example/ck_tile folder. Please check each example's subfolder.