Commit Graph

32 Commits

Author SHA1 Message Date
PoYen, Chen
aadd3ec63e Fix wrong syntax in skcheck expr 2024-08-16 10:09:46 +00:00
PoYen, Chen
a4c6029a3d Fix skcheck logic 2024-08-16 10:08:01 +00:00
PoYen, Chen
5805f5aa73 Remove group mode from appendkv kernel 2024-08-16 10:04:48 +00:00
PoYen, Chen
a8a2275aca Fix wrong arugment count 2024-08-13 08:42:23 +00:00
PoYen, Chen
a0d2163045 Remove dropout code in splitkv kernel 2024-08-08 10:21:34 +00:00
PoYen, Chen
9d9c5a6c24 Fix compilation errors 2024-08-08 08:26:55 +00:00
PoYen, Chen
291e9b4bbb Separate splitkv/non-splitkv args/traits 2024-08-08 08:07:03 +00:00
PoYen, Chen
15d0034a64 Add paged-kv codegen logic for appendkv kernels 2024-08-07 04:19:45 +00:00
PoYen, Chen
b98985262d Add missing kernel arguments for group mode 2024-08-06 14:54:07 +00:00
PoYen, Chen
8779716403 Fix uneven split checking logic 2024-08-06 01:17:14 +00:00
PoYen, Chen
381f7e90e0 Merge branch 'develop' into feature/fmha-fwd-appendkv 2024-08-04 02:12:20 +00:00
PoYen, Chen
e7969b9fd2 Add template argument 'kIsPagedKV' for splitkv kernels 2024-08-02 10:14:51 +00:00
carlushuang
b3f86e79dd workaround rocm-6.2 compiler issue (#1421) 2024-07-31 16:03:59 +08:00
PoYen, Chen
4280a07d2a Refine pipeline padding settings 2024-07-24 11:37:56 +00:00
PoYen, Chen
08b4e8a125 Fix wrong rope key for fp8 pipeline 2024-07-24 06:06:07 +00:00
PoYen, Chen
d59e098ec4 Fix wrong pipeline 2024-07-24 05:17:57 +00:00
PoYen, Chen
29c9b650b5 Align commit message to the real comment 2024-07-24 05:14:00 +00:00
PoYen, Chen
c7b7b44883 Add comment for why I just 't' for all padding flags 2024-07-24 05:13:16 +00:00
PoYen, Chen
2192bbc68a Rename RotaryEmbeddingEnum 2024-07-23 07:50:50 +00:00
PoYen, Chen
fffd6799e6 Instantiate multiple kernels for RoPE approaches 2024-07-20 02:28:21 +00:00
PoYen, Chen
391210ed9e Pass RoPE kernel args 2024-07-14 23:18:32 +00:00
PoYen, Chen
b34ddf5f71 Merge remote-tracking branch 'origin/feature/cond-add-splitkv' into feature/fmha-fwd-appendkv 2024-07-12 04:42:45 +00:00
carlushuang
8182976c37 [CK_TILE] wa prec, remove sgpr offset for inline asm (#1356)
* wa prec, remove sgpr offset for inline asm

* macro for set tile

* ignore unused param if no kernel instances in host API

* fix more prec issue

* cache buffer resource

* fix

* support pre-nop

* clear tile by vector type members

* add workaround to reduce scratch memory

* conditionally enable workaround code

* enable workaround start from certain build version

* fallback set_tile() implementation from certain build version

* undo template argument changes

* put dummy asm in load_raw()

* fix comments, refactor s_nop inside buffer_load

---------

Co-authored-by: PoYen, Chen <PoYen.Chen@amd.com>
2024-07-08 11:09:55 -07:00
PoYen, Chen
1c070380fa Merge branch 'feature/cond-add-splitkv' into feature/fmha-fwd-appendkv 2024-07-08 07:13:34 +00:00
Po Yen Chen
0cb2e06ddc [CK_TILE] fmha forward split-kv + combine kernels (#1338)
* FA fwd dropout

* FA bwd

* epilogue reuse

* CMakeLists update

* [CK_TILE] support alibi (#1269)

* add alibi support

* fix code

* update code based on comment

* Support more hdim

* fix fp8 bias

* support seqlen_k=0 case

* remove unused printf

* fix format

---------

Co-authored-by: rocking <ChunYu.Lai@amd.com>

* now fwd/bwd can build

* bwd alibi

* add bwd validation stream_config

* update generated filenames

* update bwd kernel launch

* CK_TILE_HOST_DEVICE in philox

* Transpose -> transpose

* format

* format

* format

* Generate the instance for FA required

* format

* fix error in WarpGemm

* Add num_splits option and dummy split-kv api method

* Generate fmha_fwd_splitkv()

* Add SplitKV kernel codegen logics

* Add SplitKV combine kernel codegen logics

* Fix mismatched return type

* Clean-up code

* Replace sentinel value before storing

* Fix wrong layout of LSE/LSEacc/Oacc

* Format codes

* Fix o_acc memory error

* Fix wrong kBlockSize used in policy

* Reduce # of combine kernels

* Fix split-kv combine kernel name

* Fix wrong LDS indexing logics

* Fix wrong loop counter step logic

* Undo vector size changes

* Remove no-longer used field

* Remove in-consistent comment

* Remove debug statements in example

* Remove more debug statements

* Add constness to local variables

* Clearn up generate.py

* Fix unstable clang-format comment

* Remove unused include directive

* Use shorter template parameter name

* Enable non-split-kv blobs

* Update license date

* Print num_splits conditionally

* Undo disabling data types

* Remove unnessary tile size for fp8

* Fix wrong pipeline args for fp8

* Fix example output format

* Remove more debug code in combine pipeline

* Add stride kernel arguments for LSE/O acc workspace

* Re-order split-kv pipeline call operator arguments

* Pass LSE/O strides in kernel argument

* Re-order pipeline call operator arguments

* Use tensor_descriptor to locate LSEacc elements

* Support providing invalid element for tensor view

* Set invalid element value for LSEacc tensor view

* Remove hand-written store_tile() code

* Remove necessary value-overwrite logic

* Add transposed lds descriptor

* Support load_tile() for tile_window_with_static_lengths<>

* Undo removing necessary value-overwrite logic

* Use read descriptor to locate lds elements

* Simplify pipeline source code

* Add constraint to kMaxSplits

* Default use kMaxSplits=64 in generate.py

* Revert "Add constraint to kMaxSplits"

This reverts commit 0a2132d758.

* Revert "Default use kMaxSplits=64 in generate.py"

This reverts commit c7d9c80b77.

* Decide alignment by the padding parameter

* Remove no-longer used utility functions

* Remove not-working code

* Add comment & remove no-longer used code

* Fix computation errors

* Add heuristic to override num_splits option

* Add constraint to kMaxSplits

* Fix compilation error

* Clean up pipeline code

* Wrap pointer access as lambda function

* Rename confusing methods

* Use kLogMasSplits as template parameter

* Finish splitkv combine kernel codegen

* Update kMaxSplits limit

* Use smaller kM0 for splitkv combine kernel

* Ignore droupout flag in splitkv pipeline

* Unify flag usage

* Add back flag kStoreLSE

* Merge lambda calls in pipeline

* Fix compilation errors

* Avoid all empty splits

* Always check for empty loop in splitkv pipelines

* Re-order parameters

* Remove redundant p_drop option check

* Add traits/problem for fwd splitkv kernel

* Conditionally enable uneven split boundary checks

* Add comment for the splitkv traits field

* Change even split criteria

* Re-order statements

* Refine occupancy value for hdim=128&256

* Refine occupancy value for hdim=32&64

* Remove redundant kernel argument

* Separate fmha bwd codegen logics

* Separate fmha fwd codegen logics

* Remove redundant direction parameter in fwd&bwd codegen logics

* Support generate multiple APIs for an example

* Let 'api' an alias of 'direction' option

* Remove choices for the 'direction' option

* Use dictionary to config all the functions

* Move fmha splitkv codegen logics to other file

* Add fwd_splitkv api for tile_example_fmha_fwd

---------

Co-authored-by: danyao12 <danyao12>
Co-authored-by: carlushuang <carlus.huang@amd.com>
Co-authored-by: rocking <ChunYu.Lai@amd.com>
Co-authored-by: Jing Zhang <jizhan@amd.com>
2024-06-26 17:41:15 +08:00
PoYen, Chen
4e6c28522c Fix wrong K values after appending 2024-06-25 10:12:13 +00:00
PoYen, Chen
344902732a Sync kernel name with the codegen 2024-06-24 14:50:25 +00:00
PoYen, Chen
bace0e5df0 Add init codegen logic for fmha fwd appendkv 2024-06-24 12:33:51 +00:00
PoYen, Chen
51487f238a Support generate multiple APIs for an example 2024-06-21 10:04:17 +00:00
PoYen, Chen
925d25ff47 Remove redundant direction parameter in fwd&bwd codegen logics 2024-06-21 09:41:06 +00:00
PoYen, Chen
0ef4fb2a16 Separate fmha fwd codegen logics 2024-06-21 09:22:15 +00:00
PoYen, Chen
070e89d7ad Separate fmha bwd codegen logics 2024-06-21 08:35:15 +00:00