Skip to content

Add support of a ck_tile fmha pipeline which enabling using paged-kv without splitting kv #1292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wants to merge 955 commits into
base: main
Choose a base branch
from
Draft

Add support of a ck_tile fmha pipeline which enabling using paged-kv without splitting kv #1292

wants to merge 955 commits into from
+193 −9

Conversation

qianfengz
Copy link
Contributor

@qianfengz qianfengz commented Jul 18, 2025 •
edited
Loading

CK tile implemented new pipeline to better support PagedBlockDiagonal attention-bias types.

qianfengz and others added 30 commits September 22, 2024 22:16
@qianfengz
Move ck decoder codes to xformers/csrc/attention/hip_decoder folder
21ae9d9
@qianfengz
Sync to latest ck_tile commits for fixing NaN when seqlen_k == 0
fb3628d
@qianfengz
Separate the kernel/pipeline dispatch into two files for infer/forward
ffa9906
@qianfengz
Remove unused member variable in GroupedForwardParams
221860e
@tenpercent
delete autogenerated files
6a07c16
@tenpercent
delete autogenerated files (2)
74355e9
@qianfengz
Initial add support of fmha-forward splitk (copmiling passed)
0dbdc5f
@qianfengz
Add generated files under hip_decoder into gitignore list
6b8ddde
@tenpercent
apply black and fix lint
cf9be1c
@tenpercent
rewrite hipified split-k decoder invocation to ck-tile style
08219dc
@tenpercent
Merge pull request #25 from tenpercent/refactor-hip-decoder
f37fb3d 
[Refactor] change ck decoder invocation way from old CK to CK-Tile
@qianfengz
Force kPadSeqLenQ == true for grouped mode splitkv-combine kernel traits
9c8d2f1
@qianfengz
Add compile-time checking to save compile-time
761e8a5
@tenpercent
add dockerfile
eb50024
@tenpercent
migrate base docker image to manylinux
669ee34
@tenpercent
build a wheel
0a97ed6
@tenpercent
rename dockerfile
f8129c3
@tenpercent
lint
a0221e5
@tenpercent
Merge pull request #26 from ROCm/dockerfile
9975759 
Create a dockerfile with prebuilt xformers on ROCm
@tenpercent
Try adding docker image build workflow
4d2a37d
@tenpercent
Merge branch 'develop' into dockerfile
eddb1ec
@tenpercent
add newline
ea3b796
@tenpercent
add newline
983fc19
@tenpercent
Merge pull request #27 from ROCm/dockerfile
d6ea535 
Try adding docker image build workflow
@tenpercent
Update README.md
eb986e1 
ROCm specific installation instruction
@tenpercent
Merge pull request #28 from ROCm/tenpercent-patch-1
ac0b05c 
Update README.md
@qianfengz
Remove directly including of <ck_tile/host.hpp>
e391974
@qianfengz
Synchronize with latest ck develop commit
9d03beb
@qianfengz
Remove the printing in attention_forward_generic_ck_tiled.cpp
0a4d420
@qianfengz
Tune the TilePartitioner for splitkv-combine kernel
a1c788e
qianfengz and others added 24 commits March 27, 2025 05:05
@qianfengz
Correct the condition for using merge_nhead_groups_seqlen_q
da59ab4
@qianfengz
Merge pull request #64 from ROCm/ck-fmha-enable-knob
ffe4808 
add environment knob for turning off ck fmha
@qianfengz
Merge pull request #63 from ROCm/remove-legacy-decoder
58e1103 
remove legacy ck decoder
@qianfengz
Fix to make hip_fmha compilable on torch-2.8
02e7602
@qianfengz
Merge remote-tracking branch 'upstream/main' into develop
39addc8
@qianfengz
Add support of BlockDiagonalCausalLocalAttentionPaddedKeysMask with c…
8fdfa85 
…k.FwOp
@qianfengz
Import all used attn_bias types and remove prefix when referring to t…
e8fdbba 
…he attn_bias types
@qianfengz
Remove efficient_attention_forward_decoder_ck from the interface decl…
89967ff 
…aration in attention.cpp
@qianfengz
Return logsumexp as std::optional<at::Tensor> in efficient_attention_…
a893fd5 
…forward_ck()
@tenpercent
Update ck pin (#66)
8c203d8 
* move composable kernel pin

* update trait api

* define problems and fix more traits

* fix some kargs

* modify more kargs

* more api fixes

* run clang-format

* fix python lints

* update ck pin

* update CK pin to include API fixes on CK side
@qianfengz
Update and synchronize with the latest ck_tile kernel arguments chang…
8e2050c 
…e (min_seqlen_q added)
@ltqin
add fmha grouped infer pagedkv dispatch
9e9fda3
@ltqin
limit k==kv for pagedkv
ebc7f5a
@ltqin
remove logits_soft_cap and paged limit
60860a1
@ltqin
limit seq_len_q
2ecbc3a
@qianfengz
Update to latest ck develop commit to include ck_tile PR-2405
a622744
@qianfengz
Renaming in pagedkv_dispatch
a969435
@ltqin
remove gappy constraints and set to support only pagedkv
8ab0b15
@ltqin
Synchronize to latest ck_tile(add pagedkv for large seq_len_q)
47a2681
@qianfengz
Selecting MTile (128 or 64) for calling grouped_infer_pagedkv_mask_bi…
04c4ff7 
…as_dropout_dispatch
@qianfengz
Merge pull request #68 from ROCm/ck_tile/kvcache_prefill
f62f580 
Add attention pagedkv prefill pipeline(for improving  large seq_len_q  performance)
@qianfengz
Clarify the usage of param.use_split_kv and param.use_paged_kvcache
3e13f86
@ltqin
remove selecting mtile,just use 128
25e836f
@qianfengz
Merge pull request #69 from ROCm/ck_tile/kvcache_prefill_ch_tile_size
0fa8753 
Change pagedkv kernel mtile size
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 18, 2025
@qianfengz qianfengz marked this pull request as draft July 18, 2025 08:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: rocm
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@qianfengz @facebook-github-bot @tenpercent @poyenc @johnnynunez @ltqin